Fix device handling - properly support both CPU and CUDA

app.py

@@ -102,10 +102,13 @@ def load_model():
     model = Network(cfg, src_lang, tgt_lang)
     
     # Load pretrained weights if available
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
     if os.path.exists('./LM_MODEL.pth'):
         try:
-            # Try loading with map_location for CPU compatibility
-            checkpoint = torch.load('./LM_MODEL.pth', map_location='cpu')
+            # Load with proper device mapping
+            checkpoint = torch.load('./LM_MODEL.pth', map_location=device)
             if 'state_dict' in checkpoint:
                 state_dict = checkpoint['state_dict']
             else:
@@ -122,11 +125,15 @@ def load_model():
             print(f"Warning: Could not load full model weights: {e}")
             print("Continuing with randomly initialized weights")
     
+    model = model.to(device)
     model.eval()
     return model, src_lang, tgt_lang, cfg
 
 # Process image and text
 def process_input(image, text_input, model, src_lang, tgt_lang, cfg):
+    # Get device
+    device = next(model.parameters()).device
+    
     # Transform image
     diagram_transform = T_diagram.Compose([
         T_diagram.Resize(cfg.diagram_size),
@@ -135,7 +142,7 @@ def process_input(image, text_input, model, src_lang, tgt_lang, cfg):
         T_diagram.Normalize()
     ])
     
-    diagram = diagram_transform(image).unsqueeze(0)
+    diagram = diagram_transform(image).unsqueeze(0).to(device)
     
     # Process text input
     # Create a simple text structure
@@ -159,28 +166,28 @@ def process_input(image, text_input, model, src_lang, tgt_lang, cfg):
     # Get text indices
     text_token, text_sect_tag, text_class_tag = get_text_index(combine_text, src_lang)
     
-    # Convert to tensors
+    # Convert to tensors and move to device
     text_dict = {
-        'token': torch.LongTensor([text_token]),
-        'sect_tag': torch.LongTensor([text_sect_tag]),
-        'class_tag': torch.LongTensor([text_class_tag]),
-        'len': torch.LongTensor([len(text_token)])
+        'token': torch.LongTensor([text_token]).to(device),
+        'sect_tag': torch.LongTensor([text_sect_tag]).to(device),
+        'class_tag': torch.LongTensor([text_class_tag]).to(device),
+        'len': torch.LongTensor([len(text_token)]).to(device)
     }
     
     # Get variables and arguments
     var_arg_positions, var_values, arg_values = get_var_arg(combine_text, cfg)
     
     var_dict = {
-        'pos': torch.LongTensor([var_arg_positions]),
-        'len': torch.LongTensor([len(var_arg_positions)]),
+        'pos': torch.LongTensor([var_arg_positions]).to(device),
+        'len': torch.LongTensor([len(var_arg_positions)]).to(device),
         'var_value': var_values,
         'arg_value': arg_values
     }
     
     # Create dummy expression dict for inference
     exp_dict = {
-        'exp': torch.LongTensor([[1]]),  # SOS token
-        'len': torch.LongTensor([1]),
+        'exp': torch.LongTensor([[1]]).to(device),  # SOS token
+        'len': torch.LongTensor([1]).to(device),
         'answer': 0
     }
     

core/network.py

@@ -43,8 +43,9 @@ class MLMTransformerPretrain(nn.Module):
         return transformer_outputs
 
     def load_model(self, model_path):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         pretrain_dict = torch.load(
-            model_path, map_location="cpu"
+            model_path, map_location=device
         )
         pretrain_dict_model = pretrain_dict['state_dict'] \
                                 if 'state_dict' in pretrain_dict else pretrain_dict
@@ -156,8 +157,9 @@ class Network(nn.Module):
             p.requires_grad = False
 
     def load_model(self, model_path):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         pretrain_dict = torch.load(
-            model_path, map_location="cpu"
+            model_path, map_location=device
         )
         pretrain_dict_model = pretrain_dict['state_dict'] \
                             if 'state_dict' in pretrain_dict else pretrain_dict