fix test

2026-08-11 11:24:31 +08:00 · 2019-12-01 10:02:35 +01:00
parent dff5907e48
commit a01b5be848
1 changed files with 1 additions and 1 deletions
@@ -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,