TypeError: only integer tensors of a single element can be converted to an index

[bvsn]

Got this error when initializing Unet model with the aux_params.

Exception happens In the TrainEpoch.batch_update while calling self.model.forward(x) which relates to the SegmentationModel.forward. In the SegmentationModel.forward we can see inconsistent return: masks or tuple(masks, labels).

Full Traceback:

Traceback (most recent call last):
  File "/Project/poc/ssm/run.py", line 288, in <module>
    train_logs = train_epoch.run(train_loader)
  File "/Project/poc/venv/lib/python3.9/site-packages/segmentation_models_pytorch/utils/train.py", line 51, in run
    loss, y_pred = self.batch_update(x, y)
  File "/Project/poc/venv/lib/python3.9/site-packages/segmentation_models_pytorch/utils/train.py", line 91, in batch_update
    loss = self.loss(prediction, y)
  File "/Project/poc/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Project/poc/venv/lib/python3.9/site-packages/segmentation_models_pytorch/utils/losses.py", line 36, in forward
    return 1 - F.f_score(
  File "/Project/poc/venv/lib/python3.9/site-packages/segmentation_models_pytorch/utils/functional.py", line 57, in f_score
    tp = torch.sum(gt * pr)
TypeError: only integer tensors of a single element can be converted to an index

Versions: python - 3.9 torch - 1.12.1 segmentation_models_pytorch - 0.3.0