Update tasks/Model_Loader.py

tasks/Model_Loader.py

@@ -1,7 +1,7 @@
 import torch
 
 class M5(torch.nn.Module):
-    def __init__(self, num_classes=10):
+    def __init__(self, num_classes=2):  # Ensure it matches dataset labels (chainsaw/environment)
         super(M5, self).__init__()
         self.conv1 = torch.nn.Conv1d(in_channels=1, out_channels=32, kernel_size=80, stride=4)
         self.bn1 = torch.nn.BatchNorm1d(32)
@@ -26,13 +26,17 @@ class M5(torch.nn.Module):
         x = self.fc1(x)
         return x
 
-    def load_model(model_path, num_classes=2):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model = M5(num_classes=num_classes).to(device)
-        model.load_state_dict(torch.load(model_path, map_location=device))
-        model.eval()  # Set model to evaluation mode
-        return model, device
+def load_model(model_path, num_classes=2):
+    """
+    Load trained M5 model.
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = M5(num_classes=num_classes).to(device)
+    model.load_state_dict(torch.load(model_path, map_location=device))
+    model.eval()  # Set model to evaluation mode
+    return model, device
 
 if __name__ == "__main__":
-    model, device = load_model("m5_audio_classification.pth")
-print("✅ Model successfully loaded!")
+    model, device = load_model("quantized_teacher_m5_static.pth")
+    print("✅ Model successfully loaded!")
+