Fix CPU bug, overhaul model runner, and update to lightning >=2.0

casanovo/config.py

@@ -65,7 +65,6 @@ class Config:
         train_from_scratch=bool,
         save_model=bool,
         model_save_folder_path=str,
-        save_weights_only=bool,
         every_n_train_steps=int,
         accelerator=str,
         devices=int,

casanovo/config.yaml

@@ -115,8 +115,6 @@ train_from_scratch: True
 save_model: True
 # Path to saved checkpoints
 model_save_folder_path: ""
-# Set to "False" to save the PyTorch model instance
-save_weights_only: True
 # Model validation and checkpointing frequency in training steps
 every_n_train_steps: 50_000
 # The hardware accelerator to use. Must be one of:

casanovo/denovo/model.py

@@ -114,6 +114,7 @@ def __init__(
         **kwargs: Dict,
     ):
         super().__init__()
+        self.save_hyperparameters()
 
         # Build the model.
         if custom_encoder is not None:

casanovo/denovo/model_runner.py

@@ -59,7 +59,6 @@ def __init__(
                 pl.callbacks.ModelCheckpoint(
                     dirpath=config.model_save_folder_path,
                     save_top_k=-1,
-                    save_weights_only=config.save_weights_only,
                     every_n_train_steps=config.every_n_train_steps,
                 )
             ]
@@ -247,23 +246,9 @@ def initialize_model(self, train: bool) -> None:
             )
             raise FileNotFoundError("Could not find the model weights file")
 
-        accelerator_class = str(type(self.trainer.accelerator))
-        if "CUDA" in accelerator_class:
-            map_location_device = "cuda"
-        elif "TPU" in accelerator_class:
-            map_location_device = "xla"
-        elif "HPU" in accelerator_class:
-            map_location_device = "hpu"
-        elif "IPU" in accelerator_class:
-            map_location_device = "ipu"
-        # FIXME: Handle the case for mps separately.
-        else:
-            map_location_device = "cpu"
-
-        self.model = Spec2Pep().load_from_checkpoint(
+        self.model = Spec2Pep.load_from_checkpoint(
             self.model_filename,
-            map_location=torch.device(map_location_device),
-            **model_params,
+            map_location=torch.empty(1).device,  # Use the default device.
         )
 
     def initialize_data_module(
@@ -289,7 +274,7 @@ def initialize_data_module(
             n_devices = self.trainer.num_devices
             train_bs = self.config.train_batch_size // n_devices
             eval_bs = self.config.predict_batch_size // n_devices
-        except AttributeError as err:
+        except AttributeError:
             raise RuntimeError("Please use `initialize_trainer()` first.")
 
         self.loaders = DeNovoDataModule(