Fix: CUDA availability check and warning in deeptb.py and nnsk.py (#193)

dptb/nn/deeptb.py

@@ -294,6 +294,10 @@ def from_reference(
         transform: bool = True,
         **kwargs
         ):
+        if device == 'cuda':
+            if not torch.cuda.is_available():
+                device = 'cpu'
+                log.warning("CUDA is not available. The model will be loaded on CPU.")
 
         ckpt = torch.load(checkpoint, map_location=device)
         common_options = {
@@ -441,7 +445,11 @@ def from_reference(
         ):
         # the mapping from the parameters of the ref_model and the current model can be found using
         # reference model's idp and current idp
-
+        if device == 'cuda':
+            if not torch.cuda.is_available():
+                device = 'cpu'
+                log.warning("CUDA is not available. The model will be loaded on CPU.")
+
         ckpt = torch.load(checkpoint, map_location=device)
         common_options = {
             "dtype": dtype,

dptb/nn/nnsk.py

@@ -514,6 +514,12 @@ def from_reference(
             del v1_model
 
         else:
+            if device == 'cuda':
+                if not torch.cuda.is_available():
+                    device = 'cpu'
+                    log.warning("CUDA is not available. The model will be loaded on CPU.")
+                    common_options.update({"device": device})
+
             f = torch.load(checkpoint, map_location=device)
             for k,v in common_options.items():
                 if v is None:
@@ -685,6 +691,11 @@ def _from_model_v1(
         # idp_sk.get_orbpair_soc_maps()
         idp_sk.get_sksoc_maps()
 
+        if device == 'cuda':
+            if not torch.cuda.is_available():
+                device = 'cpu'
+                log.warning("CUDA is not available. The model will be loaded on CPU.")
+
         nnsk_model = cls(basis=basis, idp_sk=idp_sk,  onsite=onsite,
                           hopping=hopping, overlap=overlap, soc=soc, std=std,freeze=freeze, push=push, dtype=dtype, device=device)