chore: ensure consistent .pt device naming

Author: ma595

era5_training/infer.py

@@ -21,7 +21,7 @@ def main():
     input_data = load_nc_dataset(args.test_data_dir / Path(prefix + "-input.nc"))
     pred_reference = load_nc_dataset(args.test_data_dir / Path(prefix + "-predict.nc"))
 
-    model_path = args.scripted_model_dir / Path(f"nlgw_{prefix}_gpu_scripted.pt")
+    model_path = args.scripted_model_dir / Path(f"nlgw_{prefix}_{device}_scripted.pt")
     print(f"loading model {model_path}...")
     model = torch.jit.load(model_path)
 

utils/function_training.py

@@ -411,7 +411,7 @@ def Inference_and_Save_AttentionUNet(
                 xdata.to_netcdf(f"test-data/unet-{k}.nc")
 
             print("scripting...")
-            script_to_torchscript(model, filename="nlgw_unet_gpu_scripted.pt")
+            script_to_torchscript(model, filename=f"nlgw_unet_{device}_scripted.pt")
             print("complete")
 
         # write to netCDF