When training on a classification problem with multiclass and using CrossEntropyLoss, the targets should integer, and hence targets.long() is needed. But when training on multilabel or binary, we may invoke BCEWithLogitsLoss, which expects your targets to be in float.
I made changes in Dataset as follows:
if CRITERION_PARAMS.train_criterion_name == "BCEWithLogitsLoss":
    # Make changes to reshape rather than in Trainer.
    y = torch.as_tensor(y, dtype=torch.float32).view(-1, 1)
else:
    y = torch.as_tensor(y, dtype=torch.long)

However, I soon realized that this does not work since DataLoader will (collate?) the shape wrongly. What I expect is the following:
original_y_shape -> (batch_size, ) or (batch_size, 1)
modified_y_shape -> (batch_size, 1) # to match shapes with BCEWithLogitsLoss

But if I change the shape inside Dataset, the DataLoader gives me (batch_size, 1, 1) instead.
One easy way is to change the shape in the training function directly. But I would like to keep the changes in Dataset unless there is a strong reason not to (speed issues).