fix device mismatch in torchscript test

pytorch · Feb 15, 2022 · 482cedd · 482cedd
1 parent 864f98e
commit 482cedd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/torchaudio_unittest/functional/torchscript_consistency_impl.py b/test/torchaudio_unittest/functional/torchscript_consistency_impl.py
@@ -654,7 +654,7 @@ def func(tensor):
 
     def test_compute_power_spectral_density_matrix_with_mask(self):
         def func(tensor):
-            mask = torch.rand(201, 100)
+            mask = torch.rand(201, 100, device=tensor.device)
             return F.compute_power_spectral_density_matrix(tensor, mask=mask)
 
         tensor = torch.rand(2, 201, 100, dtype=torch.cfloat)