Merge pull request #1388 from BenCrulis/retain_graph

Add a variable to control whether to retain the graph or not in the backward pass

avalanche/training/templates/base_sgd.py

@@ -114,6 +114,9 @@ def __init__(
         )
         """ Eval mini-batch size. """
 
+        self.retain_graph: bool = False
+        """ Retain graph when calling loss.backward(). """
+
         if evaluator is None:
             evaluator = EvaluationPlugin()
         elif callable(evaluator):
@@ -220,7 +223,7 @@ def training_epoch(self, **kwargs):
 
     def backward(self):
         """Run the backward pass."""
-        self.loss.backward()
+        self.loss.backward(retain_graph=self.retain_graph)
 
     def optimizer_step(self):
         """Execute the optimizer step (weights update)."""