Update model.py

model.py

@@ -2,9 +2,9 @@ import torch
 import torch.nn as nn
 import segmentation_models_pytorch as smp
 
-class UNetWithRegNetYBackbone(nn.Module):
+class UNet(nn.Module):
     """
-    UNet model for multi-class segmentation with RegNetY_006 backbone.
+    UNet model for multi-class segmentation.
     Designed for multi-spectral input images (e.g., 13 Sentinel-2 bands) and multiple output classes.
     """
 
@@ -21,7 +21,7 @@ class UNetWithRegNetYBackbone(nn.Module):
             num_classes (int): Number of output classes (e.g., 4 for clear, cloud types, and shadow).
             freeze_encoder (bool): If True, freezes the encoder weights during training.
         """
-        super(UNetWithRegNetYBackbone, self).__init__()
+        super(UNet, self).__init__()
 
 
         self.unet = smp.Unet(