First, I format my mask in 1 and 0 (cause when loading, values are 255 and some images noise/artefact appear) :
mask = np.where(mask>mask.mean(), 1, 0)
mask = torch.from_numpy(mask) 
mask = mask.long()

It is now torch.int64, should I parse it to float ?
Then, for predictions, I do :
out = model(inputs)
_, predictions = torch.max(out, 1)

prediction is now a torch.int64, should I parse the result in float also here ? Should I do something like :
predictions = predictions.float()
predictions.requires_grad = True

?