fix: load from cpu when saved on gpu

.gitignore

@@ -10,6 +10,7 @@ data/
 .vscode/
 *.pt
 *~
+.vscode/
 
 # Notebook to python
 forest_example.py

pytorch_tabnet/abstract_model.py

@@ -64,10 +64,10 @@ def __post_init__(self):
  self.batch_size = 1024
  self.virtual_batch_size = 128
  torch.manual_seed(self.seed)
- torch.manual_seed(self.seed)
  # Defining device
  self.device = torch.device(define_device(self.device_name))
- print(f"Device used : {self.device}")
+ if self.verbose != 0:
+ print(f"Device used : {self.device}")
 
  def __update__(self, **kwargs):
  """
@@ -330,6 +330,7 @@ def load_weights_from_unsupervised(self, unsupervised_model):
  if self.network.state_dict().get(new_param) is not None:
  # update only common layers
  update_state_dict[new_param] = weights
+
  self.network.load_state_dict(update_state_dict)
 
  def save_model(self, path):
@@ -380,6 +381,7 @@ def load_model(self, filepath):
  with zipfile.ZipFile(filepath) as z:
  with z.open("model_params.json") as f:
  loaded_params = json.load(f)
+ loaded_params["device_name"] = self.device_name
  with z.open("network.pt") as f:
  try:
  saved_state_dict = torch.load(f, map_location=self.device)
@@ -399,6 +401,7 @@ def load_model(self, filepath):
  self._set_network()
  self.network.load_state_dict(saved_state_dict)
  self.network.eval()
+
  return
 
  def _train_epoch(self, train_loader):
@@ -539,7 +542,6 @@ def _set_network(self):
  epsilon=self.epsilon,
  virtual_batch_size=self.virtual_batch_size,
  momentum=self.momentum,
- device_name=self.device_name,
  mask_type=self.mask_type,
  ).to(self.device)
 

pytorch_tabnet/pretraining.py

@@ -182,7 +182,6 @@ def _set_network(self):
  epsilon=self.epsilon,
  virtual_batch_size=self.virtual_batch_size,
  momentum=self.momentum,
- device_name=self.device_name,
  mask_type=self.mask_type,
  ).to(self.device)
 

pytorch_tabnet/tab_network.py

@@ -298,7 +298,6 @@ def __init__(
  epsilon=1e-15,
  virtual_batch_size=128,
  momentum=0.02,
- device_name="auto",
  mask_type="sparsemax",
  ):
  super(TabNetPretraining, self).__init__()
@@ -499,7 +498,6 @@ def __init__(
  epsilon=1e-15,
  virtual_batch_size=128,
  momentum=0.02,
- device_name="auto",
  mask_type="sparsemax",
  ):
  """
@@ -538,7 +536,6 @@ def __init__(
  Batch size for Ghost Batch Normalization
  momentum : float
  Float value between 0 and 1 which will be used for momentum in all batch norm
- device_name : {'auto', 'cuda', 'cpu'}
  mask_type : str
  Either "sparsemax" or "entmax" : this is the masking function to use
  """
@@ -581,15 +578,6 @@ def __init__(
  mask_type,
  )
 
- # Defining device
- if device_name == "auto":
- if torch.cuda.is_available():
- device_name = "cuda"
- else:
- device_name = "cpu"
- self.device = torch.device(device_name)
- self.to(self.device)
-
  def forward(self, x):
  x = self.embedder(x)
  return self.tabnet(x)

pytorch_tabnet/utils.py

@@ -311,5 +311,7 @@ def define_device(device_name):
  return "cuda"
  else:
  return "cpu"
+ elif device_name == "cuda" and not torch.cuda.is_available():
+ return "cpu"
  else:
  return device_name