Update model.py

model.py

@@ -56,7 +56,7 @@ def load_model(version='c', device='cpu'):
         repo_type="model"
     )
 
-    model = SimpleCNN(model_type=model_type, num_classes=4).to(device)
+    model = SimpleCNN(model_type=model_type, num_classes=6).to(device)
     state_dict = torch.load(weights_path, map_location=device)
     model.load_state_dict(state_dict, strict=False)
     model.eval()