Fix PyTorch 2.6 weights_only loading and syntax warnings

app.py

@@ -25,7 +25,10 @@ checkpoint_path = os.path.join("checkpoints", "BIT_LEVIR", "best_ckpt.pt")
 if not os.path.exists(checkpoint_path):
     raise FileNotFoundError(f"Checkpoint not found at: {checkpoint_path}")
 
-checkpoint = torch.load(checkpoint_path, map_location=device)
+try:
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+except TypeError:
+    checkpoint = torch.load(checkpoint_path, map_location=device)
 model.load_state_dict(checkpoint["model_G_state_dict"])
 model.to(device)
 model.eval()

models/basic_model.py

@@ -28,9 +28,12 @@ class CDEvaluator():
     def load_checkpoint(self, checkpoint_name='best_ckpt.pt'):
 
         if os.path.exists(os.path.join(self.checkpoint_dir, checkpoint_name)):
-            # load the entire checkpoint
-            checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
-                                    map_location=self.device)
+            try:
+                checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
+                                        map_location=self.device, weights_only=False)
+            except TypeError:
+                checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
+                                        map_location=self.device)
 
             self.net_G.load_state_dict(checkpoint['model_G_state_dict'])
             self.net_G.to(self.device)

models/evaluator.py

@@ -64,8 +64,12 @@ class CDEvaluator():
 
         if os.path.exists(os.path.join(self.checkpoint_dir, checkpoint_name)):
             self.logger.write('loading last checkpoint...\n')
-            # load the entire checkpoint
-            checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name), map_location=self.device)
+            try:
+                checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
+                                        map_location=self.device, weights_only=False)
+            except TypeError:
+                checkpoint = torch.load(os.path.join(self.checkpoint_dir, checkpoint_name),
+                                        map_location=self.device)
 
             self.net_G.load_state_dict(checkpoint['model_G_state_dict'])
 

models/networks.py

@@ -263,7 +263,7 @@ class BASE_Transformer(ResNet):
         mlp_dim = 2*dim
 
         self.with_pos = with_pos
-        if with_pos is 'learned':
+        if with_pos == 'learned':
             self.pos_embedding = nn.Parameter(torch.randn(1, self.token_len*2, 32))
         decoder_pos_size = 256//4
         self.with_decoder_pos = with_decoder_pos
@@ -294,9 +294,9 @@ class BASE_Transformer(ResNet):
 
     def _forward_reshape_tokens(self, x):
         # b,c,h,w = x.shape
-        if self.pool_mode is 'max':
+        if self.pool_mode == 'max':
             x = F.adaptive_max_pool2d(x, [self.pooling_size, self.pooling_size])
-        elif self.pool_mode is 'ave':
+        elif self.pool_mode == 'ave':
             x = F.adaptive_avg_pool2d(x, [self.pooling_size, self.pooling_size])
         else:
             x = x