Repository: graphnet-team/graphnet
Issue: Model.load_from_checkpoint fails

**Describe the bug**
The pytorch_lightning method `.load_from_checkpoint` fails because it attempts to access `model.load_state_dict` with kwargs that the underlying `torch.nn.Module` requires but our identically named method doesn't accept.

**To Reproduce**
Steps to reproduce the behavior:
1. Go to main and run

```
from torch.optim.adam import Adam
from graphnet.models import StandardModel
from graphnet.models.detector.icecube import IceCubeKaggle
from graphnet.models.gnn import DynEdge
from graphnet.models.graph_builders import KNNGraphBuilder
from graphnet.models.task.reconstruction import DirectionReconstructionWithKappa
from graphnet.training.callbacks import PiecewiseLinearLR
from graphnet.training.loss_functions import VonMisesFisher3DLoss

from typing import Any, Dict, List, Optional

def build_model_and_load_checkpoint(target: str,
                checkpoint_path: str) -> StandardModel:
    """Builds GNN """

    # Building Fresh Model
    detector = IceCubeKaggle(
        graph_builder=KNNGraphBuilder(nb_nearest_neighbours=8),
    )
    gnn = DynEdge(
        nb_inputs=detector.nb_outputs,
        global_pooling_schemes=["min", "max", "mean"],
    )

    task = DirectionReconstructionWithKappa(
            hidden_size=gnn.nb_outputs,
            target_labels=target,
            loss_function=VonMisesFisher3DLoss())

    model = StandardModel(
        detector=detector,
        gnn=gnn,
        tasks=[task],
        optimizer_class=Adam,
        optimizer_kwargs={"lr": 1e-03, "eps": 1e-03},
    )

    #Load In Some Checkpoint
    model = StandardModel.load_from_checkpoint(detector = detector,
                                               gnn = gnn,
                                               optimizer_class = Adam,
                                               tasks = [task],
                                               optimizer_kwargs = {"lr": 1e-03, "eps": 1e-03},
                                               checkpoint_path= checkpoint_path)
    return model

checkpoint_path = '/home/iwsatlas1/oersoe/phd/kaggle_competition/example_code/checkpoints/epoch=49-step=1581650.ckpt'

model_from_checkpoint = build_model_and_load_checkpoint(target = 'direction',
                                                        checkpoint_path = checkpoint_path)
```

2. see error

**Expected behavior**
Should not fail.

**Full traceback**
Please include the full error message to allow for debugging
```
Singularity> python load_model_checkpoint.py 
/home/iwsatlas1/oersoe/.local/lib/python3.9/site-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 4.0.0-unsupported is an invalid version and will not be supported in a future release
  warnings.warn(
/home/iwsatlas1/oersoe/.local/lib/python3.9/site-packages/pkg_resources/__init__.py:116: PkgResourcesDeprecationWarning: 4.0.0-unsupported is an invalid version and will not be supported in a future release
  warnings.warn(
graphnet [MainProcess] INFO     2023-03-15 15:13:21 - KNNGraphBuilder.__init__ - Writing log to logs/graphnet_20230315-151321.log
Traceback (most recent call last):
  File "/home/iwsatlas1/oersoe/phd/tmp/kaggle_debug/load_model_checkpoint.py", line 49, in <module>
    model_from_checkpoint = build_model_and_load_checkpoint(target = 'direction',
  File "/home/iwsatlas1/oersoe/phd/tmp/kaggle_debug/load_model_checkpoint.py", line 39, in build_model_and_load_checkpoint
    model = StandardModel.load_from_checkpoint(detector = detector,
  File "/home/iwsatlas1/oersoe/.local/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 137, in load_from_checkpoint
    return _load_from_checkpoint(
  File "/home/iwsatlas1/oersoe/.local/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 205, in _load_from_checkpoint
    return _load_state(cls, checkpoint, strict=strict, **kwargs)
  File "/home/iwsatlas1/oersoe/.local/lib/python3.9/site-packages/pytorch_lightning/core/saving.py", line 259, in _load_state
    keys = obj.load_state_dict(checkpoint["state_dict"], strict=strict)
TypeError: load_state_dict() got an unexpected keyword argument 'strict'
```



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

diff --git a/src/graphnet/models/model.py b/src/graphnet/models/model.py
index cafef6a6..74afeed7 100644
--- a/src/graphnet/models/model.py
+++ b/src/graphnet/models/model.py
@@ -246,14 +246,14 @@ class Model(Logger, Configurable, LightningModule, ABC):
         self.info(f"Model state_dict saved to {path}")
 
     def load_state_dict(
-        self, path: Union[str, Dict]
+        self, path: Union[str, Dict], **kargs: Optional[Any]
     ) -> "Model":  # pylint: disable=arguments-differ
         """Load model `state_dict` from `path`."""
         if isinstance(path, str):
             state_dict = torch.load(path)
         else:
             state_dict = path
-        return super().load_state_dict(state_dict)
+        return super().load_state_dict(state_dict, **kargs)
 
     @classmethod
     def from_config(  # type: ignore[override]