Fix: mask dtype → LongTensor for DiceLoss one_hot compatibility

train.py

@@ -145,8 +145,8 @@ class ESCDDDataset(Dataset):
         # Apply augmentations
         if self.transform:
             augmented = self.transform(image=img, mask=mask)
-            img = augmented["image"]    # (1, H, W) float tensor
-            mask = augmented["mask"]    # (H, W) long tensor
+            img = augmented["image"]           # (1, H, W) float tensor
+            mask = augmented["mask"].long()    # (H, W) LongTensor — required by DiceLoss one_hot
         else:
             img = torch.from_numpy(img).unsqueeze(0) / 255.0
             mask = torch.from_numpy(mask).long()
@@ -331,7 +331,7 @@ def train():
 
         for batch_idx, (images, masks) in enumerate(train_loader):
             images = images.to(device)
-            masks = masks.to(device)
+            masks = masks.to(device).long()   # Ensure LongTensor on GPU
 
             optimizer.zero_grad()
 
@@ -358,7 +358,7 @@ def train():
         with torch.no_grad():
             for images, masks in val_loader:
                 images = images.to(device)
-                masks = masks.to(device)
+                masks = masks.to(device).long()   # Ensure LongTensor on GPU
 
                 with torch.amp.autocast(device_type="cuda", enabled=cfg.USE_AMP):
                     logits = model(images)
@@ -387,7 +387,7 @@ def train():
             all_per_class = {name: [] for name in cfg.CLASS_NAMES}
             with torch.no_grad():
                 for images, masks in val_loader:
-                    images, masks = images.to(device), masks.to(device)
+                    images, masks = images.to(device), masks.to(device).long()
                     with torch.amp.autocast(device_type="cuda", enabled=cfg.USE_AMP):
                         logits = model(images)
                     m = compute_metrics(logits, masks, cfg.NUM_CLASSES)