log model

mllam · Oct 7, 2024 · de27e9a · de27e9a
1 parent b53bab5
commit de27e9a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py
@@ -48,7 +48,7 @@ def log_image(self, key, images):
 
     def log_model(self, model):
         # Create model signature
-        #signature = infer_signature(X.numpy(), model(X).detach().numpy())
+        #signature = infer_signature(train_dataset.numpy(), model(train_dataset).detach().numpy())
         mlflow.pytorch.log_model(model, "model")
 
 
@@ -361,7 +361,7 @@ def main(input_args=None):
 
         # Log the model
         training_logger.log_model(model)
-
+        # data_module.train_dataloader().dataset.data
 
 if __name__ == "__main__":
     main()