diff --git a/NN/models/advanced_transformer_trading.py b/NN/models/advanced_transformer_trading.py
index 23bd7be..f06ece4 100644
--- a/NN/models/advanced_transformer_trading.py
+++ b/NN/models/advanced_transformer_trading.py
@@ -1182,9 +1182,9 @@ class TradingTransformerTrainer:
                     price_data_1h=batch.get('price_data_1h'),
                     price_data_1d=batch.get('price_data_1d'),
                     btc_data_1m=batch.get('btc_data_1m'),
-                    cob_data=batch['cob_data'],
-                    tech_data=batch['tech_data'],
-                    market_data=batch['market_data'],
+                    cob_data=batch.get('cob_data'),  # Use .get() to handle missing key
+                    tech_data=batch.get('tech_data'),
+                    market_data=batch.get('market_data'),
                     position_state=batch.get('position_state'),
                     price_data=batch.get('price_data')  # Legacy fallback
                 )