Update model/Model_RGB.py

, map_location=torch.device('cpu')

model/Model_RGB.py

@@ -324,7 +324,7 @@ class EstimateRGB(nn.Module):
 
     def _make_model(self):
         model = SAR(self.cfg.backbone, in_channels=self.cfg.in_channels).to(self.cfg.device)
-        checkpoint = torch.load(self.cfg.checkpoint)
+        checkpoint = torch.load(self.cfg.checkpoint, map_location=torch.device('cpu'))
         if 'net' in checkpoint:
             model.load_state_dict(checkpoint['net'])
         else:
@@ -351,7 +351,7 @@ class EstimateRGB(nn.Module):
         import onnxruntime as ort
         print('export begin')
         model = SAR(self.cfg.backbone, in_channels=self.cfg.in_channels).to(self.cfg.device)
-        checkpoint = torch.load(self.cfg.checkpoint)
+        checkpoint = torch.load(self.cfg.checkpoint, map_location=torch.device('cpu'))
         if 'net' in checkpoint:
             model.load_state_dict(checkpoint['net'])
         else: