Hi Frank , thank you for your reply here is my function, where I call the loss function:
    def batch_update_cross(self, x, y):
        self.optimizer.zero_grad()
        prediction = self.model.forward(x)
        depth = y.shape[1]
        merged_mask = torch.zeros([2, 928 ,928])
        for j in range(depth):
            merged_mask[y[:, j, :, :] == 1] = int(j + 1)
        y = merged_mask.long().cuda()
        loss = self.loss(prediction, y)
        loss.backward()
        self.optimizer.step()
        return loss, prediction

I first tried it with :
y = merged_mask.int().cuda()

but then I got this error:
RuntimeError: Expected object of scalar type Long but got scalar type Int for argument #2 'target' in call to _thnn_nll_loss2d_forward

So i changed .int() into .long() and this following error appears:
C:/w/1/s/tmp_conda_3.7_100118/conda/conda-bld/pytorch_1579082551706/work/aten/src/THCUNN/SpatialClassNLLCriterion.cu:104: block: [9,0,0], thread: [526,0,0] Assertion `t >= 0 && t < n_classes` failed.
Traceback (most recent call last):
  File "C:/Users/lukas/Documents/SchneidbiegeRoboter/PyTorch/Pytroch_train_model.py", line 150, in <module>
    m, o = train(train_idx=training_idx, val_idx=validation_idx, retrain=settings.retrain)
  File "C:/Users/lukas/Documents/SchneidbiegeRoboter/PyTorch/Pytroch_train_model.py", line 91, in train
    train_logs = train_epoch.run(train_loader)
  File "C:\Users\lukas\anaconda3\envs\Python\lib\site-packages\segmentation_models_pytorch\utils\train.py", line 47, in run
    loss, y_pred = self.batch_update_cross(x, y)
  File "C:\Users\lukas\anaconda3\envs\Python\lib\site-packages\segmentation_models_pytorch\utils\train.py", line 94, in batch_update_cross
    loss.backward()
  File "C:\Users\lukas\anaconda3\envs\Python\lib\site-packages\torch\tensor.py", line 195, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "C:\Users\lukas\anaconda3\envs\Python\lib\site-packages\torch\autograd\__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED

now there could be one issue I’m thinking of, my class labels aren’t 0 to 4, they are 1 to 5 does this make any difference ?
Cheers Lukas