change num classes

eth-easl · Sep 6, 2024 · b01f3bc · b01f3bc
1 parent a0be9a7
commit b01f3bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/...ts/trainer_server/internal/trainer/remote_downsamplers/test_remote_gradnorm_downsample.py b/...ts/trainer_server/internal/trainer/remote_downsamplers/test_remote_gradnorm_downsample.py
@@ -22,7 +22,7 @@ def test_sample_shape_ce(dummy_system_config: ModynConfig):
     )
     with torch.inference_mode(mode=(not sampler.requires_grad)):
         data = torch.randn(8, 10)
-        target = torch.randint(3, size=(8,))
+        target = torch.randint(2, size=(8,))
         ids = list(range(8))
         forward_outputs = model(data)
         sampler.inform_samples(ids, data, forward_outputs, target)