Hi ptrblck,
Thanks so much for your response.
I tried to change the loss_fct part but it gives me error:
AttributeError: 'tuple' object has no attribute 'device'

The error message shows up with model_train function.
This is a function I define that includes train/validate/test. I can see from the progress bar that the training part can be done without any issues. The problem starts with the validation part – which is calling this line:
ave_valid_loss, valid_f1_table, cm_table, f1_score = model_validate(model=model, data_loader=valid_data_loader)

And model_validate function is defined as below:
def model_validate(model, data_loader):
    model.eval()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)

    label_prop = data_loader.dataset.dataset.label_prop()

    total_valid_loss = 0

    batch_size = data_loader.batch_size
    num_batch = len(data_loader)

    y_pred, y_true = [], []

    # Evaluate data
    for step, batch in tqdm(enumerate(data_loader), desc="Validation...", total=num_batch):
        b_review_input_ids = batch["review_input_ids"].to(device)
        b_review_attention_mask = batch["review_attention_mask"].to(device)
        b_review_token_type_ids = batch["review_token_type_ids"].to(device)
        b_agent_input_ids = batch["agent_input_ids"].to(device)
        b_agent_attention_mask = batch["agent_attention_mask"].to(device)
        b_agent_token_type_ids = batch["agent_token_type_ids"].to(device)

        b_binarized_label = batch["binarized_label"].to(device)

        with torch.no_grad():
            (loss, logits,) = model(review_input_ids=b_review_input_ids,
                                    review_attention_mask=b_review_attention_mask,
                                    review_token_type_ids=b_review_token_type_ids,
                                    agent_input_ids=b_agent_input_ids,
                                    agent_attention_mask=b_agent_attention_mask,
                                    agent_token_type_ids=b_agent_token_type_ids,

                                    labels=b_binarized_label)

        total_valid_loss += loss.item()

        transfored_logits = torch.sigmoid(logits)
        numpy_probas = transfored_logits.detach().cpu().numpy()
        y_pred.extend(numpy_probas)
        y_true.extend(b_binarized_label.cpu().numpy())
        
    # End of an epoch of validation

    # put model to train mode again.
    model.train()

    ave_loss = total_valid_loss / (num_batch * batch_size)

    y_pred = np.array(y_pred)
    y_pred[y_pred < 0.5] = 0
    y_pred[y_pred >= 0.5] = 1
    
    # Below is in case the input and target are not the same data format
    y_pred = np.array(y_pred, dtype=np.bool)
    y_true = np.array(y_true, dtype=np.bool)
    
    # compute the various f1 score for each label
    report = classification_report(y_true, y_pred, output_dict=True)
    metrics_df = pd.DataFrame(report).transpose()
    metrics_df = metrics_df.sort_index()

    weighted_f1_score = metrics_df.loc['weighted avg', 'f1-score']
    averaged_f1_score = metrics_df.loc['macro avg', 'f1-score']

    # Calculate confusion matrix
    tn, fp, fn, tp  = confusion_matrix(y_true, y_pred).ravel()
    cm_df = pd.DataFrame(columns = ['Predicted No', 'Predicted Yes'],  
                   index = ['Actual No', 'Actual Yes']) 

    cm_df.loc['Actual No'] = [tn,fp] 
    cm_df.loc['Actual Yes'] = [fn,tp]

    return ave_loss, metrics_df, cm_df,{
        "weighted": weighted_f1_score,
        "averaged": averaged_f1_score,
    }

The traceback:
RuntimeError                              Traceback (most recent call last)
<ipython-input-14-7ca45d5cff01> in <module>
      7 model_train(model=model, train_data_loader=data_loader["train"], valid_data_loader=data_loader["valid"],
      8             test_data_loader=data_loader["test"], optimizer=optimizer, scheduler=scheduler,
----> 9             num_epochs=args.num_epochs, seed=args.seed, logger=logger, out_dir=out_dir)

~/SageMaker/Anecdotes/train_valid_test.py in model_train(model, train_data_loader, valid_data_loader, test_data_loader, logger, optimizer, scheduler, num_epochs, seed, out_dir)
    162         t_valid = time.time()
    163         model.eval()
--> 164         ave_valid_loss, valid_f1_table, cm_table, f1_score = model_validate(model=model, data_loader=valid_data_loader)
    165         # Measure how long this epoch took.
    166         validation_time = format_time(time.time() - t_valid)

~/SageMaker/Anecdotes/train_valid_test.py in model_validate(model, data_loader)
    280                                     agent_token_type_ids=b_agent_token_type_ids,
    281 
--> 282                                     labels=b_binarized_label)
    283 
    284         if num_gpus > 1:

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py in forward(self, *inputs, **kwargs)
    153             return self.module(*inputs[0], **kwargs[0])
    154         replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
--> 155         outputs = self.parallel_apply(replicas, inputs, kwargs)
    156         return self.gather(outputs, self.output_device)
    157 

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py in parallel_apply(self, replicas, inputs, kwargs)
    163 
    164     def parallel_apply(self, replicas, inputs, kwargs):
--> 165         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
    166 
    167     def gather(self, outputs, output_device):

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py in parallel_apply(modules, inputs, kwargs_tup, devices)
     83         output = results[i]
     84         if isinstance(output, ExceptionWrapper):
---> 85             output.reraise()
     86         outputs.append(output)
     87     return outputs

~/anaconda3/envs/python3/lib/python3.6/site-packages/torch/_utils.py in reraise(self)
    393             # (https://bugs.python.org/issue2651), so we work around it.
    394             msg = KeyErrorMessage(msg)
--> 395         raise self.exc_type(msg)

Thank you again for your help!