Upload continuous_learning_session.py

continuous_learning_session.py

@@ -1,6 +1,11 @@
 import random
 import logging
 import os
+import gc
+
+# Optimize CUDA memory allocation to reduce fragmentation
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -84,7 +89,8 @@ class ContinuousLearningSession:
              # If it's a python list
              adapter_params = [p for layer in self.model.flux_layers for p in layer.parameters()]
              
-        self.optimizer = optim.Adam(controller_params + adapter_params, lr=1e-4) # Reduced from 1e-3 for stability
+        # Switch to SGD to save memory (Adam uses 2x states, causing OOM on T4)
+        self.optimizer = optim.SGD(controller_params + adapter_params, lr=1e-3, momentum=0.9)
         
         self.model.train()  # Enable gradients for Controller/Adapters
         
@@ -270,6 +276,10 @@ class ContinuousLearningSession:
         # 1. Add new knowledge to Buffer
         self.replay_buffer.add(concept_id, user_input, correct_answer)
         
+        # Force cleanup before training to prevent OOM
+        gc.collect()
+        torch.cuda.empty_cache()
+
         # 2. Training Loop (Micro-Epochs)
         steps = 50 # Increase back to 50 since we have more data now!