below is the code:
weight = [3.17]
class_weight = torch.FloatTensor(weight)
loss_function = nn.BCEWithLogitsLoss(pos_weight=class_weight)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

y_pred = model(categorical_train_data, numerical_train_data)
print(y_pred.shape, train_outputs.shape)
single_loss = loss_function(y_pred, train_outputs)
#single_loss = loss_function(y_pred, torch.max(train_outputs,1)[1])

output:
torch.Size([39074, 2]) torch.Size([39074, 1])
  File "/apps/miniconda3/envs/torch/lib/python3.8/site-packages/torch/nn/functional.py", line 2124, in binary_cross_entropy_with_logits
raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
 ValueError: Target size (torch.Size([39074, 1])) must be the same as input size (torch.Size([39074, 2]))

Just checking what to do