Fix a possible crash when using checkpoints

Author: Aethor

README.md

@@ -120,6 +120,7 @@ The following parameters can be set (taken from `./tibert/run_train.py` config f
 | `segment_size`               | `128`               |
 | `encoder`                    | `"bert-base-cased"` |
 | `out_model_dir`              | `"~/tibert/model"`  |
+| `checkpoint`                 | `None`              |
 
 
 One can monitor training metrics by adding run observers using command line flags - see `sacred` documentation for more details.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tibert"
-version = "0.2.1"
+version = "0.2.2"
 description = "BERT for Coreference Resolution"
 authors = ["Arthur Amalvy <[email protected]>"]
 license = "GPL-3.0-only"

tibert/train.py

@@ -51,7 +51,7 @@ def load_train_checkpoint(
 
     model_config = config_class(**checkpoint["model_config"])
     model = model_class(model_config)
-    model.load_state_dict(checkpoint["model"])
+    model.load_state_dict(checkpoint["model"], strict=False)
 
     optimizer = torch.optim.AdamW(
         [
@@ -68,6 +68,17 @@ def load_train_checkpoint(
     return model, optimizer
 
 
+def _optimizer_to_(
+    optimizer: torch.optim.AdamW, device: torch.device
+) -> torch.optim.AdamW:
+    """From https://github.com/pytorch/pytorch/issues/2830"""
+    for state in optimizer.state.values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor):
+                state[k] = v.cuda()
+    return optimizer
+
+
 def train_coref_model(
     model: Union[BertForCoreferenceResolution, CamembertForCoreferenceResolution],
     dataset: CoreferenceDataset,
@@ -161,8 +172,8 @@ def train_coref_model(
             ],
             lr=task_lr,
         )
+    optimizer = _optimizer_to_(optimizer, device)
 
-    # Best model saving
     # -----------------
     best_f1 = 0
     best_model = model