Extend test by a bit, fix docstring

malfet · malfet · commit d15dc6b28f5f · 2024-05-21T13:40:45.000-07:00
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -7490,9 +7490,12 @@ def test_sobolengine_fast_forward_scrambled(self):
         self.test_sobolengine_fast_forward(scramble=True)
 
     def test_sobolengine_default_dtype(self):
+        engine = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123456)
+        # Check that default dtype is correctly handled
+        self.assertEqual(engine.draw(n=5).dtype, torch.float32)
         with set_default_dtype(torch.float64):
             engine = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123456)
-            # Check that default dtype is correctly handled
+            # Check that default dtype is correctly handled (when set to float64)
             self.assertEqual(engine.draw(n=5).dtype, torch.float64)
             # Check that explicitly passed dtype is adhered to
             self.assertEqual(engine.draw(n=5, dtype=torch.float32).dtype, torch.float32)
diff --git a/torch/quasirandom.py b/torch/quasirandom.py
@@ -119,7 +119,7 @@ def draw_base2(self, m: int, out: Optional[torch.Tensor] = None,
             out (Tensor, optional): The output tensor
             dtype (:class:`torch.dtype`, optional): the desired data type of the
                                                     returned tensor.
-                                                    Default: ``torch.float32``
+                                                    Default: ``None``
         """
         n = 2 ** m
         total_n = self.num_generated + n
