Update core/train_eval.py

core/train_eval.py

@@ -154,6 +154,14 @@ def train_and_evaluate(
         train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
         test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
 
+        # Debug DataLoader output
+        for batch_X, batch_y in train_loader:
+            logging.debug(f"DataLoader train batch: X_type={type(batch_X)}, X_shape={batch_X.shape}, y_type={type(batch_y)}, y_shape={batch_y.shape}")
+            break
+        for batch_X, batch_y in test_loader:
+            logging.debug(f"DataLoader test batch: X_type={type(batch_X)}, X_shape={batch_X.shape}, y_type={type(batch_y)}, y_shape={batch_y.shape}")
+            break
+
         input_size = X.shape[2]
         model = model_cls(input_size=input_size, hidden_size=hidden, num_layers=layers, output_size=horizon, dropout=dropout).to(device)
         logging.debug(f"Model initialized: {model_cls.__name__}, input_size={input_size}, hidden={hidden}, layers={layers}")
@@ -187,8 +195,12 @@ def train_and_evaluate(
                 logging.debug(f"Training Batch_X type: {type(batch_X)}, shape: {batch_X.shape}")
                 logging.debug(f"Training Batch_Y type: {type(batch_y)}, shape: {batch_y.shape}")
                 optimizer.zero_grad()
-                outputs = model(batch_X)
-                logging.debug(f"Model output shape: {outputs.shape}")
+                try:
+                    outputs = model(batch_X)
+                    logging.debug(f"Model output shape: {outputs.shape}")
+                except Exception as e:
+                    logging.error(f"Model forward error: {str(e)}, batch_X_shape={batch_X.shape}")
+                    raise
                 loss = criterion(outputs, batch_y)
                 loss.backward()
                 optimizer.step()
@@ -202,8 +214,12 @@ def train_and_evaluate(
                 for batch_X, batch_y in test_loader:
                     logging.debug(f"Validation Batch_X type: {type(batch_X)}, shape: {batch_X.shape}")
                     logging.debug(f"Validation Batch_Y type: {type(batch_y)}, shape: {batch_y.shape}")
-                    outputs = model(batch_X)
-                    logging.debug(f"Validation model output shape: {outputs.shape}")
+                    try:
+                        outputs = model(batch_X)
+                        logging.debug(f"Validation model output shape: {outputs.shape}")
+                    except Exception as e:
+                        logging.error(f"Validation model forward error: {str(e)}, batch_X_shape={batch_X.shape}")
+                        raise
                     loss = criterion(outputs, batch_y)
                     val_loss += loss.item() * batch_X.size(0)
                 val_loss /= len(test_loader.dataset)