Add AMD GPU detection, replace LSTM with MLP, and fix signal generation logic

- Detect AMD GPUs via ROCm device name in config
- Replace single-timestep LSTM in GNN model with leaner post-GNN MLP
- Pass edge_attr through AMDGATConv propagate and use ones for self-loops
- Fix live trading to sell existing positions on negative signals instead
  of skipping them entirely
- Use per-file try/except in data pipeline and batch SQLite inserts
- Import torch directly in backtester instead of dynamic __import__
- Update AMP autocast import for PyTorch 2.0+ compatibility

config.py

@@ -7,6 +7,17 @@ from dotenv import load_dotenv
 load_dotenv()
 
 
+def _check_amd_gpu() -> bool:
+    """Detect AMD GPU via ROCm device name."""
+    if not torch.cuda.is_available():
+        return False
+    try:
+        name = torch.cuda.get_device_name(0)
+        return any(kw in name for kw in ("AMD", "Radeon", "gfx"))
+    except Exception:
+        return False
+
+
 class Config:
     # Project settings
     PROJECT_NAME = "StockGNN_R9700"
@@ -20,12 +31,6 @@ class Config:
     EXTERNAL_DATA_DIR = os.path.join(DATA_DIR, "external")
     MODEL_DIR = os.path.join(BASE_DIR, "models")
 
-    # Ensure directories exist
-    os.makedirs(RAW_DATA_DIR, exist_ok=True)
-    os.makedirs(PROCESSED_DATA_DIR, exist_ok=True)
-    os.makedirs(EXTERNAL_DATA_DIR, exist_ok=True)
-    os.makedirs(MODEL_DIR, exist_ok=True)
-
     # Stock universe settings
     INITIAL_TICKERS = [
         "AAPL",
@@ -71,7 +76,7 @@ class Config:
 
     # AMD GPU settings (Radeon R9700 AI Pro - 32GB)
     DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-    AMD_GPU = True
+    AMD_GPU = _check_amd_gpu()
     GPU_MEMORY_LIMIT = 0.9  # Use 90% of 32GB = 28.8GB
     ROCM_OPT_LEVEL = "O2"  # Optimization level for ROCm ('O0', 'O1', 'O2')
     PIN_MEMORY = True  # Enable pinned memory for faster data transfer
@@ -195,5 +200,11 @@ class Config:
     EXECUTION_TIME_HORIZON = "5min"  # Time to complete execution
     MARKET_IMPACT_MODEL = "kyle"  # 'kyle', 'almgren_chriss', or 'none'
 
+    def __init__(self):
+        os.makedirs(self.RAW_DATA_DIR, exist_ok=True)
+        os.makedirs(self.PROCESSED_DATA_DIR, exist_ok=True)
+        os.makedirs(self.EXTERNAL_DATA_DIR, exist_ok=True)
+        os.makedirs(self.MODEL_DIR, exist_ok=True)
+
 
 config = Config()

live_trading.py

@@ -376,60 +376,50 @@ class LiveTradingSystem:
                 if not current_price:
                     continue
 
-                # Skip if we already have a position in this stock
-                if ticker in self.trader.current_positions:
-                    continue
+                held_qty = self.trader.current_positions.get(ticker, 0)
 
-                # Check holding period constraints
-                if self.trader._check_holding_period(ticker, timestamp):
-                    continue
-
-                # Generate signals based on prediction
-                if prediction > 0.002:  # Buy signal
-                    # Calculate position size
-                    position_size = self.trader._calculate_position_size(
-                        ticker, current_price
-                    )
-
-                    if position_size > 0:
-                        # Create buy order
+                if held_qty > 0:
+                    # Existing position: sell on negative signal (holding-period check still applies)
+                    if prediction < -0.002 and not self.trader._check_holding_period(
+                        ticker, timestamp
+                    ):
                         order = {
                             "ticker": ticker,
-                            "action": "buy",
-                            "quantity": position_size,
+                            "action": "sell",
+                            "quantity": held_qty,
                             "price": current_price,
                             "timestamp": timestamp,
                             "type": "market",
                         }
-
-                        # Submit order
                         order_id = self.broker.submit_order(order)
                         if order_id:
                             self.trader.pending_orders[order_id] = order
                             logger.info(
-                                f"Submitted buy order for {position_size} shares of {ticker} at {current_price}"
+                                f"Submitted sell order for {held_qty} shares of {ticker} at {current_price}"
                             )
-
-                elif (
-                    prediction < -0.002 and ticker in self.trader.current_positions
-                ):  # Sell signal
-                    # Create sell order
-                    order = {
-                        "ticker": ticker,
-                        "action": "sell",
-                        "quantity": self.trader.current_positions[ticker],
-                        "price": current_price,
-                        "timestamp": timestamp,
-                        "type": "market",
-                    }
-
-                    # Submit order
-                    order_id = self.broker.submit_order(order)
-                    if order_id:
-                        self.trader.pending_orders[order_id] = order
-                        logger.info(
-                            f"Submitted sell order for {self.trader.current_positions[ticker]} shares of {ticker} at {current_price}"
+                else:
+                    # No position: buy on positive signal
+                    if prediction > 0.002 and not self.trader._check_holding_period(
+                        ticker, timestamp
+                    ):
+                        position_size = self.trader._calculate_position_size(
+                            ticker, current_price
                         )
+                        if position_size > 0:
+                            order = {
+                                "ticker": ticker,
+                                "action": "buy",
+                                "quantity": position_size,
+                                "price": current_price,
+                                "timestamp": timestamp,
+                                "type": "market",
+                            }
+                            order_id = self.broker.submit_order(order)
+                            if order_id:
+                                self.trader.pending_orders[order_id] = order
+                                logger.info(
+                                    f"Submitted buy order for {position_size} shares of {ticker} at {current_price}"
+                                )
 
         except Exception as e:
             logger.error(f"Error generating trading signals: {str(e)}", exc_info=True)

src/amd/optimizations.py

@@ -290,8 +290,8 @@ class AMDGATConv(MessagePassing):
         alpha_dst = (x_dst * self.att_dst).sum(dim=-1)
         alpha = (alpha_src, alpha_dst)
 
-        # Propagate
-        out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, size=size)
+        # Propagate — pass edge_attr so correlation weights reach message()
+        out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size)
 
         # Concatenate or average heads
         if self.concat:
@@ -312,7 +312,8 @@ class AMDGATConv(MessagePassing):
         loop_index = loop_index.unsqueeze(0).repeat(2, 1)
 
         if edge_attr is not None:
-            loop_attr = edge_attr.new_zeros((num_nodes,) + edge_attr.size()[1:])
+            # Self-loop weight = 1.0 so self-messages are not zeroed out
+            loop_attr = edge_attr.new_ones((num_nodes,) + edge_attr.size()[1:])
             edge_attr = torch.cat([edge_attr, loop_attr], dim=0)
 
         edge_index = torch.cat([edge_index, loop_index], dim=1)
@@ -321,15 +322,21 @@ class AMDGATConv(MessagePassing):
     def propagate(self, edge_index, size=None, **kwargs):
         return super().propagate(edge_index, size=size, **kwargs)
 
-    def message(self, x_j, alpha_j, alpha_i, index, ptr, size_i):
+    def message(self, x_j, alpha_j, alpha_i, edge_attr, index, ptr, size_i):
         # Compute attention weights
         alpha = alpha_j + alpha_i
         alpha = F.leaky_relu(alpha, negative_slope=0.2)
         alpha = self._softmax(alpha, index, ptr, size_i)
         alpha = F.dropout(alpha, p=self.dropout, training=self.training)
 
-        # Weighted sum of values
-        return x_j * alpha.unsqueeze(-1)
+        # (E, heads, out_channels) scaled by attention
+        msg = x_j * alpha.unsqueeze(-1)
+
+        # Scale by edge correlation weight: (E, 1) → broadcasts to (E, heads, out_channels)
+        if edge_attr is not None:
+            msg = msg * edge_attr.unsqueeze(1)
+
+        return msg
 
     def _softmax(self, src, index, ptr, num_nodes):
         # Memory-efficient softmax

src/data/pipeline.py

@@ -147,25 +147,26 @@ class StockDataPipeline:
 
     def _load_data(self):
         """Load existing data from disk with memory management"""
-        try:
-            # Check memory before loading
-            if not self.memory_manager.ensure_memory(2 * 1024**3):  # 2GB
-                logger.warning("Skipping data load due to memory constraints")
-                return
+        if not self.memory_manager.ensure_memory(2 * 1024**3):  # 2GB
+            logger.warning("Skipping data load due to memory constraints")
+            return
 
-            self.price_data = self._load_pickle("price_data.pkl")
-            self.corporate_actions = self._load_pickle("corporate_actions.pkl")
-            self.sector_data = self._load_pickle("sector_data.pkl")
-            self.index_composition = self._load_pickle("index_composition.pkl")
+        for attr, filename in [
+            ("price_data", "price_data.pkl"),
+            ("corporate_actions", "corporate_actions.pkl"),
+            ("sector_data", "sector_data.pkl"),
+            ("index_composition", "index_composition.pkl"),
+        ]:
+            try:
+                setattr(self, attr, self._load_pickle(filename))
+                logger.info(f"Loaded {filename}")
+            except FileNotFoundError:
+                logger.info(f"{filename} not found, using empty store")
+            except Exception as e:
+                logger.error(f"Error loading {filename}: {e}", exc_info=True)
+                self.memory_manager.empty_cache()
 
-            logger.info("Loaded existing data from disk")
-            self.memory_manager.log_memory_usage("[After Data Load]")
-
-        except FileNotFoundError:
-            logger.info("No existing data found. Starting with empty databases.")
-        except Exception as e:
-            logger.error(f"Error loading data: {str(e)}", exc_info=True)
-            self.memory_manager.empty_cache()
+        self.memory_manager.log_memory_usage("[After Data Load]")
 
     def _save_data(self):
         """Save data to disk with memory management"""
@@ -497,19 +498,23 @@ class StockDataPipeline:
 
     def _store_index_composition(self):
         """Store index composition in the database"""
-        for index_ticker, composition in self.index_composition.items():
-            for date, members in composition.items():
-                for member in members:
-                    with sqlite3.connect(self.db_path) as conn:
-                        cursor = conn.cursor()
-                        cursor.execute(
-                            """
-                            INSERT OR REPLACE INTO index_composition
-                            (index_ticker, date, member_ticker)
-                            VALUES (?, ?, ?)
-                        """,
-                            (index_ticker, date, member),
-                        )
+        rows = [
+            (index_ticker, date, member)
+            for index_ticker, composition in self.index_composition.items()
+            for date, members in composition.items()
+            for member in members
+        ]
+        if not rows:
+            return
+        with sqlite3.connect(self.db_path) as conn:
+            conn.executemany(
+                """
+                INSERT OR REPLACE INTO index_composition
+                (index_ticker, date, member_ticker)
+                VALUES (?, ?, ?)
+                """,
+                rows,
+            )
 
     def update_alternative_data(self, tickers: List[str]):
         """
@@ -806,6 +811,20 @@ class StockDataPipeline:
                 self.memory_manager.empty_cache()
                 continue
 
+            # Compute once per day — stock universe and PIT data don't change intraday
+            pd_date = pd.Timestamp(date)
+            current_tickers = [
+                t for t in tickers
+                if t in self.price_data
+                and not self.price_data[t].empty
+                and self.price_data[t].index[0] <= pd_date <= self.price_data[t].index[-1]
+            ]
+            if not current_tickers:
+                continue
+
+            date_dt = datetime.strptime(date, "%Y-%m-%d")
+            pit_cache = {t: self.get_point_in_time_data(t, date_dt) for t in current_tickers}
+
             # Get all timestamps for this trading day
             timestamps = generate_intraday_timestamps(date)
 
@@ -813,20 +832,6 @@ class StockDataPipeline:
                 current_timestamp = timestamps[i]
                 sequence_start = timestamps[i - config.SEQUENCE_LENGTH]
 
-                # Get current universe of stocks
-                current_tickers = []
-                for ticker in tickers:
-                    if ticker in self.price_data and not self.price_data[ticker].empty:
-                        if (
-                            date >= self.price_data[ticker].index[0]
-                            and date <= self.price_data[ticker].index[-1]
-                        ):
-                            current_tickers.append(ticker)
-
-                # Skip if no stocks available
-                if not current_tickers:
-                    continue
-
                 # Create node features for each stock in the sequence
                 sequence_features = []
                 valid_tickers = []
@@ -905,13 +910,9 @@ class StockDataPipeline:
                                 continue
 
                             try:
-                                # Get sector relationship
-                                pit1 = self.get_point_in_time_data(
-                                    ticker1, datetime.strptime(date, "%Y-%m-%d")
-                                )
-                                pit2 = self.get_point_in_time_data(
-                                    ticker2, datetime.strptime(date, "%Y-%m-%d")
-                                )
+                                # Use pre-fetched PIT data — avoids O(n²) repeated calls
+                                pit1 = pit_cache[ticker1]
+                                pit2 = pit_cache[ticker2]
 
                                 if (
                                     pit1["sector"]

src/evaluation/backtester.py

@@ -7,6 +7,7 @@ from datetime import timedelta
 from typing import Dict, List, Optional, Tuple
 
 import pandas as pd
+import torch
 
 from config import config
 from src.models.trainer import GNNTrainer
@@ -61,9 +62,11 @@ class GNNBacktester:
             # Move data to the same device as the model
             data_device = data.to(config.DEVICE)
 
-            with __import__("torch").no_grad():
+            with torch.no_grad():
                 predictions = self.model(data_device)
 
+            portfolio_val = self._portfolio_value(date)
+
             for i, ticker in enumerate(data.tickers):
                 pred = predictions[i].item()
                 price = self._get_current_price(ticker, date)
@@ -71,18 +74,40 @@ class GNNBacktester:
                 if price is None or price <= 0:
                     continue
 
-                if pred > 0.002:
-                    order = {
-                        "ticker": ticker,
-                        "action": "buy",
-                        "quantity": 100,
-                        "price": price,
-                        "timestamp": str(date),
-                        "type": "market",
-                    }
-                    order_id = self.broker.submit_order(order)
-                    if order_id:
-                        trade_log.append({**order, "order_id": order_id})
+                held_qty = self.broker.positions.get(ticker, 0)
+
+                if held_qty > 0:
+                    # Existing position: sell on negative signal
+                    if pred < -0.002:
+                        order = {
+                            "ticker": ticker,
+                            "action": "sell",
+                            "quantity": held_qty,
+                            "price": price,
+                            "timestamp": str(date),
+                            "type": "market",
+                        }
+                        order_id = self.broker.submit_order(order)
+                        if order_id:
+                            trade_log.append({**order, "order_id": order_id})
+                else:
+                    # No position: buy on positive signal
+                    if pred > 0.002:
+                        max_spend = portfolio_val * config.MAX_POSITION_SIZE
+                        quantity = int(max_spend / price)
+                        if quantity <= 0:
+                            continue
+                        order = {
+                            "ticker": ticker,
+                            "action": "buy",
+                            "quantity": quantity,
+                            "price": price,
+                            "timestamp": str(date),
+                            "type": "market",
+                        }
+                        order_id = self.broker.submit_order(order)
+                        if order_id:
+                            trade_log.append({**order, "order_id": order_id})
 
             portfolio_values.append(self._portfolio_value(date))
             dates.append(date)

src/models/gnn_model.py

@@ -208,124 +208,82 @@ class CorporateActionAwareGNN(nn.Module):
             nn.Softmax(dim=1),
         )
 
-        # LSTM for temporal dependencies with AMD optimizations
-        self.lstm = nn.LSTM(
-            input_size=config.HIDDEN_CHANNELS,
-            hidden_size=config.HIDDEN_CHANNELS,
-            num_layers=2,
-            batch_first=True,
-            dropout=config.DROPOUT if config.NUM_HEADS > 1 else 0,
-            bidirectional=False,
+        # MLP applied after GNN (replaces single-timestep LSTM which was equivalent
+        # to a linear layer but 4× more parameters)
+        self.post_gnn_mlp = nn.Sequential(
+            nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS),
+            nn.SiLU(),
+            nn.LayerNorm(config.HIDDEN_CHANNELS),
         )
 
-        # Initialize LSTM weights
-        self._init_lstm_weights()
-
-        # Final prediction layer
-        self.linear = nn.Linear(
-            config.HIDDEN_CHANNELS * 2, 1
-        )  # *2 for concatenating GNN output and attention features
+        # Final prediction: GNN output (HIDDEN) concat with weighted modal sum (HIDDEN)
+        self.linear = nn.Linear(config.HIDDEN_CHANNELS * 2, 1)
 
         # Initialize weights
         self._init_weights()
 
     def _init_weights(self):
         """Initialize weights with AMD-friendly initialization"""
-        # Initialize linear layer weights
         nn.init.xavier_uniform_(self.linear.weight)
         nn.init.zeros_(self.linear.bias)
 
-        # Initialize attention layers
-        for layer in self.alternative_data_attention:
-            if isinstance(layer, nn.Linear):
-                nn.init.xavier_uniform_(layer.weight)
-                nn.init.zeros_(layer.bias)
+        for module in (self.alternative_data_attention, self.post_gnn_mlp):
+            for layer in module:
+                if isinstance(layer, nn.Linear):
+                    nn.init.xavier_uniform_(layer.weight)
+                    nn.init.zeros_(layer.bias)
 
-    def _init_lstm_weights(self):
-        """Initialize LSTM weights with orthogonal initialization for better convergence"""
-        for name, param in self.lstm.named_parameters():
-            if "weight_ih" in name:
-                nn.init.orthogonal_(param)
-            elif "weight_hh" in name:
-                nn.init.orthogonal_(param)
-            elif "bias" in name:
-                nn.init.zeros_(param)
-                # Set forget gate bias to 1 for better gradient flow
-                n = param.size(0)
-                param.data[n // 4 : n // 2].fill_(1)
+    def _split_features(self, x):
+        """Split node feature tensor into (price, news, social, corporate_action)."""
+        num_price = 5
+        num_news = len(config.NEWS_FEATURES)
+        num_social = len(config.SOCIAL_FEATURES)
+        return (
+            x[:, :, :num_price],
+            x[:, :, num_price : num_price + num_news],
+            x[:, :, num_price + num_news : num_price + num_news + num_social],
+            x[:, :, -1:],
+        )
 
     def forward(self, data):
         x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
 
-        # Split features into different types
-        num_price_features = 5
-        num_news_features = len(config.NEWS_FEATURES)
-        num_social_features = len(config.SOCIAL_FEATURES)
+        price_features, news_features, social_features, corporate_action_flags = (
+            self._split_features(x)
+        )
 
-        # Price features (first 5)
-        price_features = x[:, :, :num_price_features]
         price_processed = self.price_processor(price_features)
-
-        # News features (next num_news_features)
-        news_start = num_price_features
-        news_end = news_start + num_news_features
-        news_features = x[:, :, news_start:news_end]
         news_processed = self.news_processor(news_features)
-
-        # Social features (next num_social_features)
-        social_start = news_end
-        social_end = social_start + num_social_features
-        social_features = x[:, :, social_start:social_end]
         social_processed = self.social_processor(social_features)
-
-        # Corporate action flag (last feature)
-        corporate_action_flags = x[:, :, -1:]
         corporate_action_features = self.corporate_action_mlp(corporate_action_flags)
 
-        # Apply temporal attention to each feature type
         price_attended = self.temporal_attention(price_processed)
         news_attended = self.temporal_attention(news_processed)
         social_attended = self.temporal_attention(social_processed)
         corporate_attended = self.temporal_attention(corporate_action_features)
 
-        # Calculate attention weights for alternative data
         alternative_features = torch.cat(
             [price_attended, news_attended, social_attended], dim=1
         )
-
         attention_weights = self.alternative_data_attention(alternative_features)
 
-        # Apply attention weights
         weighted_price = price_attended * attention_weights[:, 0].unsqueeze(1)
         weighted_news = news_attended * attention_weights[:, 1].unsqueeze(1)
         weighted_social = social_attended * attention_weights[:, 2].unsqueeze(1)
 
-        # Combine features
         combined_features = torch.cat(
             [weighted_price, weighted_news, weighted_social, corporate_attended], dim=1
         )
 
-        # Process through GNN with gradient checkpointing for memory efficiency
         x = self._gnn_forward(combined_features, edge_index, edge_attr)
+        x = self.post_gnn_mlp(x)  # (num_stocks, HIDDEN_CHANNELS)
 
-        # Process through LSTM for temporal dependencies
-        lstm_input = x.unsqueeze(1)  # (num_stocks, 1, hidden_size)
-        lstm_out, _ = self.lstm(lstm_input)
-        lstm_out = lstm_out.squeeze(1)
+        # Weighted combination of attended modalities: (N, HIDDEN_CHANNELS)
+        modal_combined = weighted_price + weighted_news + weighted_social
 
-        # Combine GNN output with attention features
-        attention_features = torch.cat(
-            [
-                attention_weights[:, 0].unsqueeze(1),
-                attention_weights[:, 1].unsqueeze(1),
-                attention_weights[:, 2].unsqueeze(1),
-            ],
-            dim=1,
-        )
+        # Concat GNN output with modal summary → (N, HIDDEN_CHANNELS * 2)
+        x = torch.cat([x, modal_combined], dim=1)
 
-        x = torch.cat([lstm_out, attention_features], dim=1)
-
-        # Final prediction
         return self.linear(x)
 
     def _gnn_forward(self, x, edge_index, edge_attr):
@@ -345,39 +303,14 @@ class CorporateActionAwareGNN(nn.Module):
 
     def get_attention_weights(self, data):
         """Get attention weights for interpretability"""
-        x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
+        x = data.x
+        price_features, news_features, social_features, _ = self._split_features(x)
 
-        # Process through the model up to attention
-        num_price_features = 5
-        num_news_features = len(config.NEWS_FEATURES)
-        num_social_features = len(config.SOCIAL_FEATURES)
-
-        price_features = x[:, :, :num_price_features]
-        news_features = x[
-            :, :, num_price_features : num_price_features + num_news_features
-        ]
-        social_features = x[
-            :,
-            :,
-            num_price_features + num_news_features : num_price_features
-            + num_news_features
-            + num_social_features,
-        ]
-        corporate_action_flags = x[:, :, -1:]
-
-        price_processed = self.price_processor(price_features)
-        news_processed = self.news_processor(news_features)
-        social_processed = self.social_processor(social_features)
-        corporate_action_features = self.corporate_action_mlp(corporate_action_flags)
-
-        price_attended = self.temporal_attention(price_processed)
-        news_attended = self.temporal_attention(news_processed)
-        social_attended = self.temporal_attention(social_processed)
+        price_attended = self.temporal_attention(self.price_processor(price_features))
+        news_attended = self.temporal_attention(self.news_processor(news_features))
+        social_attended = self.temporal_attention(self.social_processor(social_features))
 
         alternative_features = torch.cat(
             [price_attended, news_attended, social_attended], dim=1
         )
-
-        attention_weights = self.alternative_data_attention(alternative_features)
-
-        return attention_weights
+        return self.alternative_data_attention(alternative_features)

src/models/trainer.py

@@ -6,7 +6,8 @@ from typing import Dict, List, Tuple
 import numpy as np
 import torch
 import torch.nn as nn
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import autocast
+from torch.cuda.amp import GradScaler
 from torch_geometric.loader import DataLoader
 from tqdm import tqdm
 
@@ -92,6 +93,7 @@ class GNNTrainer:
             # Training
             self.model.train()
             epoch_train_loss = 0.0
+            num_train_batches = 0
             start_time = time.time()
 
             for batch in tqdm(
@@ -110,6 +112,7 @@ class GNNTrainer:
 
                     # Mixed precision training
                     with autocast(
+                        config.DEVICE,
                         enabled=config.MIXED_PRECISION,
                         dtype=self.amd_optimizer.get_precision_dtype(),
                     ):
@@ -130,6 +133,7 @@ class GNNTrainer:
                     self.scaler.update()
 
                     epoch_train_loss += loss.item()
+                    num_train_batches += 1
 
                     # Memory management
                     self.memory_manager.auto_manage_memory(threshold=0.8)
@@ -139,12 +143,13 @@ class GNNTrainer:
                     self.memory_manager.empty_cache()
                     continue
 
-            epoch_train_loss /= len(train_loader)
+            epoch_train_loss = epoch_train_loss / num_train_batches if num_train_batches else 0.0
             self.train_losses.append(epoch_train_loss)
 
             # Validation
             self.model.eval()
             epoch_val_loss = 0.0
+            num_val_batches = 0
 
             with torch.no_grad():
                 for batch in val_loader:
@@ -159,6 +164,7 @@ class GNNTrainer:
                         batch = batch.to(self.device, non_blocking=config.PIN_MEMORY)
 
                         with autocast(
+                            config.DEVICE,
                             enabled=config.MIXED_PRECISION,
                             dtype=self.amd_optimizer.get_precision_dtype(),
                         ):
@@ -166,6 +172,7 @@ class GNNTrainer:
                             loss = self.criterion(out, batch.y)
 
                         epoch_val_loss += loss.item()
+                        num_val_batches += 1
 
                         # Memory management
                         self.memory_manager.auto_manage_memory(threshold=0.8)
@@ -178,7 +185,7 @@ class GNNTrainer:
                         self.memory_manager.empty_cache()
                         continue
 
-            epoch_val_loss /= len(val_loader)
+            epoch_val_loss = epoch_val_loss / num_val_batches if num_val_batches else 0.0
             self.val_losses.append(epoch_val_loss)
 
             # Update learning rate scheduler