✅ Improve dropout cvg.

ENSTA-U2IS-AI · Oct 25, 2023 · a01c46c · a01c46c
1 parent aa3c47b
commit a01c46c
Showing 1 changed file with 24 additions and 9 deletions.
diff --git a/tests/baselines/test_mc_dropout.py b/tests/baselines/test_mc_dropout.py
@@ -28,9 +28,9 @@ def test_standard(self):
         )
         summary(net)
 
-        _ = net.criterion
-        _ = net.configure_optimizers()
-        _ = net(torch.rand(1, 3, 32, 32))
+        net.criterion
+        net.configure_optimizers()
+        net(torch.rand(1, 3, 32, 32))
 
 
 class TestStandardWideBaseline:
@@ -49,9 +49,9 @@ def test_standard(self):
         )
         summary(net)
 
-        _ = net.criterion
-        _ = net.configure_optimizers()
-        _ = net(torch.rand(1, 3, 32, 32))
+        net.criterion
+        net.configure_optimizers()
+        net(torch.rand(1, 3, 32, 32))
 
 
 class TestStandardVGGBaseline:
@@ -67,9 +67,24 @@ def test_standard(self):
             num_estimators=4,
             arch=11,
             groups=1,
+            last_layer_dropout=True,
         )
         summary(net)
 
-        _ = net.criterion
-        _ = net.configure_optimizers()
-        _ = net(torch.rand(1, 3, 32, 32))
+        net.criterion
+        net.configure_optimizers()
+        net(torch.rand(1, 3, 32, 32))
+
+        net = VGG(
+            num_classes=10,
+            in_channels=3,
+            loss=nn.CrossEntropyLoss,
+            optimization_procedure=optim_cifar10_resnet18,
+            version="mc-dropout",
+            num_estimators=4,
+            arch=11,
+            groups=1,
+            last_layer_dropout=True,
+        )
+        net.eval()
+        net(torch.rand(1, 3, 32, 32))