fix icarl test

AlbinSou · Jan 25, 2024 · 26f13ed · 26f13ed
1 parent 0875745
commit 26f13ed
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/tests/training/test_strategies.py b/tests/training/test_strategies.py
@@ -853,13 +853,14 @@ def test_expertgate(self):
             strategy.train(experience)
 
     def test_icarl(self):
-        model, optimizer, criterion, benchmark = self.init_scenario(multi_task=False)
-        model = SimpleMLP(input_size=6, hidden_size=10)
-        optimizer = SGD(model.parameters(), lr=0.000001)
+        _, optimizer, criterion, benchmark = self.init_scenario(multi_task=False)
+        fe = torch.nn.Linear(6, 10)
+        cls = torch.nn.Sigmoid()
+        optimizer = SGD([*fe.parameters(), *cls.parameters()], lr=0.001)
 
         strategy = ICaRL(
-            model.features,
-            model.classifier,
+            fe,
+            cls,
             optimizer,
             20,
             buffer_transform=None,