Repository: graphnet-team/graphnet
Issue: 'StandardModel' object has no attribute '_inference_trainer'

**Describe the bug**
When you load in a saved state dict to a freshly built model, you cannot do inference with that model because the inference code itself relies on `Trainer` which is only available if `model.fit()` has been called. 

**To Reproduce**
Steps to reproduce the behavior:
1. Run any training example
2. Replace the code in the example that trains the model, with the line
```
model.load_state_dict(state_dict_path)
```
where state_dict_path is the path of a saved state dict.
5. Proceed to run the part of the example code that does inference with this model.
6. See error

**Expected behavior**
A clear and concise description of what you expected to happen.

**Full traceback**
Please include the full error message to allow for debugging
```
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_931/847469493.py in <module>
      1 # Inference
----> 2 results = inference(model, config)
      3 
      4 # Read predictions from disk (fastest)
      5 #results = pd.read_csv('')

/tmp/ipykernel_931/157537399.py in inference(model, config)
    156         test_dataloader,
    157         prediction_columns=model.prediction_columns,
--> 158         additional_attributes=model.additional_attributes,
    159     )
    160     # Save predictions and model to file

/kaggle/working/software/graphnet/src/graphnet/models/model.py in predict_as_dataframe(self, dataloader, prediction_columns, node_level, additional_attributes, index_column)
    138             )
    139 
--> 140         predictions_torch = self.predict(dataloader)
    141         predictions = (
    142             torch.cat(predictions_torch, dim=1).detach().cpu().numpy()

/kaggle/working/software/graphnet/src/graphnet/models/standard_model.py in predict(self, dataloader)
    166         """Return predictions for `dataloader`."""
    167         self.inference()
--> 168         return super().predict(dataloader)

/kaggle/working/software/graphnet/src/graphnet/models/model.py in predict(self, dataloader)
     93         self.train(mode=False)
     94 
---> 95         predictions_list = self._inference_trainer.predict(self, dataloader)
     96         assert len(predictions_list), "Got no predictions"
     97 

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
   1184                 return modules[name]
   1185         raise AttributeError("'{}' object has no attribute '{}'".format(
-> 1186             type(self).__name__, name))
   1187 
   1188     def __setattr__(self, name: str, value: Union[Tensor, 'Module']) -> None:

AttributeError: 'StandardModel' object has no attribute '_inference_trainer'
```

**Additional context**
Add any other context about the problem here.


Changed files:
src/graphnet/models/model.py, src/graphnet/models/standard_model.py

diff --git a/src/graphnet/models/model.py b/src/graphnet/models/model.py
index fac5b1be..0d667b62 100644
--- a/src/graphnet/models/model.py
+++ b/src/graphnet/models/model.py
@@ -27,11 +27,8 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
     def forward(self, x: Union[Tensor, Data]) -> Union[Tensor, Data]:
         """Forward pass."""
 
-    def fit(
+    def _construct_trainers(
         self,
-        train_dataloader: DataLoader,
-        val_dataloader: Optional[DataLoader] = None,
-        *,
         max_epochs: int = 10,
         gpus: Optional[Union[List[int], int]] = None,
         callbacks: Optional[List[Callback]] = None,
@@ -42,8 +39,6 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
         distribution_strategy: Optional[str] = "ddp",
         **trainer_kwargs: Any,
     ) -> None:
-        """Fit `Model` using `pytorch_lightning.Trainer`."""
-        self.train(mode=True)
 
         if gpus:
             accelerator = "gpu"
@@ -52,7 +47,7 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
             accelerator = "cpu"
             devices = None
 
-        trainer = Trainer(
+        self._trainer = Trainer(
             accelerator=accelerator,
             devices=devices,
             max_epochs=max_epochs,
@@ -77,21 +72,65 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
             **trainer_kwargs,
         )
 
+    def fit(
+        self,
+        train_dataloader: DataLoader,
+        val_dataloader: Optional[DataLoader] = None,
+        *,
+        max_epochs: Optional[int] = 10,
+        gpus: Optional[Union[List[int], int]] = None,
+        callbacks: Optional[List[Callback]] = None,
+        ckpt_path: Optional[str] = None,
+        logger: Optional[Logger] = None,
+        log_every_n_steps: Optional[int] = 1,
+        gradient_clip_val: Optional[float] = None,
+        distribution_strategy: Optional[str] = "ddp",
+        **trainer_kwargs: Any,
+    ) -> None:
+        """Fit `Model` using `pytorch_lightning.Trainer`."""
+        self.train(mode=True)
+
+        self._construct_trainer(
+            max_epochs=max_epochs,
+            gpus=gpus,
+            callbacks=callbacks,
+            ckpt_path=ckpt_path,
+            logger=logger,
+            log_every_n_steps=log_every_n_steps,
+            gradient_clip_val=gradient_clip_val,
+            distribution_strategy=distribution_strategy,
+            **trainer_kwargs,
+        )
+
         try:
-            trainer.fit(
+            self._trainer.fit(
                 self, train_dataloader, val_dataloader, ckpt_path=ckpt_path
             )
         except KeyboardInterrupt:
             self.warning("[ctrl+c] Exiting gracefully.")
             pass
 
-    def predict(self, dataloader: DataLoader) -> List[Tensor]:
+    def predict(
+        self,
+        dataloader: DataLoader,
+        gpus: Optional[Union[List[int], int]] = None,
+    ) -> List[Tensor]:
         """Return predictions for `dataloader`.
 
         Returns a list of Tensors, one for each model output.
         """
         self.train(mode=False)
 
+        if not hasattr(self, "_inference_trainer"):
+            self._construct_trainer(
+                gpus=gpus,
+            )
+        elif gpus is not None:
+            self.warning(
+                "A `Trainer` instance has already been constructed, possibly "
+                "when the model was trained. Will use this to get predictions. "
+                f"Argument `gpus = {gpus}` will be ignored."
+            )
         predictions_list = self._inference_trainer.predict(self, dataloader)
         assert len(predictions_list), "Got no predictions"
 
@@ -111,6 +150,7 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
         node_level: bool = False,
         additional_attributes: Optional[List[str]] = None,
         index_column: str = "event_no",
+        gpus: Optional[Union[List[int], int]] = None,
     ) -> pd.DataFrame:
         """Return predictions for `dataloader` as a DataFrame.
 
@@ -136,8 +176,10 @@ class Model(Configurable, LightningModule, LoggerMixin, ABC):
                 "doesn't resample batches; or do not request "
                 "`additional_attributes`."
             )
-
-        predictions_torch = self.predict(dataloader)
+        predictions_torch = self.predict(
+            dataloader=dataloader,
+            gpus=gpus,
+        )
         predictions = (
             torch.cat(predictions_torch, dim=1).detach().cpu().numpy()
         )
diff --git a/src/graphnet/models/standard_model.py b/src/graphnet/models/standard_model.py
index 1863eef2..8d778880 100644
--- a/src/graphnet/models/standard_model.py
+++ b/src/graphnet/models/standard_model.py
@@ -162,7 +162,14 @@ class StandardModel(Model):
                 task.train_eval()
         return self
 
-    def predict(self, dataloader: DataLoader) -> List[Tensor]:
+    def predict(
+        self,
+        dataloader: DataLoader,
+        gpus: Optional[Union[List[int], int]] = None,
+    ) -> List[Tensor]:
         """Return predictions for `dataloader`."""
         self.inference()
-        return super().predict(dataloader)
+        return super().predict(
+            dataloader=dataloader,
+            gpus=gpus,
+        )