Upload folder using huggingface_hub

MLBaseModelDriver.py

@@ -64,7 +64,10 @@ class MLBaseModelDriver:
 
         model = model_class(input_dim=4)
         state_dict = torch.load(model_file, weights_only=False)
-        model.load_state_dict(state_dict)
+        checkpoint = torch.load(model_path, map_location=torch.device("cpu"))
+        model.load_state_dict(checkpoint['model_state_dict'])
+        self.preprocessor = checkpoint.get('preprocessor', None)
+        self.input_dim = checkpoint.get('input_dim', None)
         model.eval()
 
         with open(scaler_file, 'rb') as f: