Greatly reduce trainning mem consuption

Corentin-Allaire · Dec 7, 2023 · 791498d · 791498d
1 parent 397cc2e
commit 791498d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/Examples/Scripts/Python/MLAmbiguityResolution/train_ambiguity_solver.py b/Examples/Scripts/Python/MLAmbiguityResolution/train_ambiguity_solver.py
@@ -188,7 +188,7 @@ def scoringBatch(batch: list[pd.DataFrame], Optimiser=0) -> tuple[int, int, floa
         nb_part += 1
         # Normalise the loss to the batch size
         batch_loss = batch_loss / len(b_data[0])
-        loss += batch_loss
+        loss += batch_loss.item()
         # Perform the gradient descend if an optimiser was specified
         if Optimiser:
             batch_loss.backward()

diff --git a/Examples/Scripts/Python/MLAmbiguityResolution/train_seed_solver.py b/Examples/Scripts/Python/MLAmbiguityResolution/train_seed_solver.py
@@ -216,7 +216,7 @@ def scoringBatch(batch: list[pd.DataFrame], Optimiser=0) -> tuple[int, int, floa
         nb_part += 1
         # Normalise the loss to the batch size
         batch_loss = batch_loss / len(b_data[0])
-        loss += batch_loss
+        loss += batch_loss.item()
         # Perform the gradient descent if an optimiser was specified
         if Optimiser:
             batch_loss.backward()