diff --git a/test/model/deepar/test_auxillary_outputs.py b/test/model/deepar/test_auxillary_outputs.py
index 0e73b29..8f1bb68 100644
--- a/test/model/deepar/test_auxillary_outputs.py
+++ b/test/model/deepar/test_auxillary_outputs.py
@@ -62,7 +62,7 @@ def test_distribution():
             *[data_entry[k] for k in input_names]
         )
 
-        assert distr.sample(num_samples).shape == (
+        assert distr.sample(torch.Size([num_samples])).shape == (
             num_samples,
             batch_size,
             seq_len,