address comments

test/test_utils.py

@@ -10,15 +10,20 @@
 @pytest.mark.parametrize(
     "tensors",
     [
-        list(map(torch.tensor, [[1, 2, 3], [5, 6]])),
-        list(map(torch.tensor, [[2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0, 10.0]])),
+        [
+            torch.tensor([2.0, 3.0], dtype=dtype),
+            torch.tensor([4.0, 5.0, 6.0], dtype=dtype),
+            torch.tensor([7.0, 8.0, 9.0, 10.0], dtype=dtype),
+        ]
+        for dtype in [torch.int, torch.float16, torch.float32]
     ],
 )
 def test_pad_and_stack(tensors: list):
     stacked_and_padded = left_pad_and_stack_1D(tensors)
+
+    assert stacked_and_padded.dtype == torch.float32
     assert stacked_and_padded.shape == (len(tensors), max(len(t) for t in tensors))
 
     ref = torch.concat(tensors).to(dtype=stacked_and_padded.dtype)
 
     assert torch.sum(torch.nan_to_num(stacked_and_padded, nan=0)) == torch.sum(ref)
-    assert torch.nanmean(stacked_and_padded) == torch.nanmean(ref)