refactor: code refactoring;

pypots/data/saving/pickle.py

@@ -32,12 +32,13 @@ def pickle_dump(data: object, path: str) -> None:
         create_dir_if_not_exist(extract_parent_dir(path))
         with open(path, "wb") as f:
             pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+        logger.info(f"Successfully saved to {path}")
     except Exception as e:
         logger.error(
-            f"❌ Pickling failed. No cache data saved. Please investigate the error below.\n{e}"
+            f"❌ Pickling failed. No cache data saved. Investigate the error below:\n{e}"
         )
-        return None
-    logger.info(f"Successfully saved to {path}")
+
+    return None
 
 
 def pickle_load(path: str) -> object:
@@ -58,6 +59,7 @@ def pickle_load(path: str) -> object:
         with open(path, "rb") as f:
             data = pickle.load(f)
     except Exception as e:
-        logger.error(f"❌ Loading data failed. Operation aborted. See info below:\n{e}")
+        logger.error(f"❌ Loading data failed. Operation aborted. Investigate the error below:\n{e}"
         return None
+
     return data

pypots/nn/modules/etsformer/layers.py

@@ -160,8 +160,9 @@ def forward(self, x):
             f = fft.rfftfreq(t)[self.low_freq :]
 
         x_freq, index_tuple = self.topk_freq(x_freq)
-        f = repeat(f, "f -> b f d", b=x_freq.size(0), d=x_freq.size(2)).to(x_freq.device)
-        f = rearrange(f[index_tuple], "b f d -> b f () d").to(x_freq.device)
+        device = x_freq.device
+        f = repeat(f, "f -> b f d", b=x_freq.size(0), d=x_freq.size(2)).to(device)
+        f = rearrange(f[index_tuple], "b f d -> b f () d").to(device)
 
         return self.extrapolate(x_freq, f, t)