Update tests/utils/test_data_type_utils.py

Co-authored-by: Sukriti Sharma <[email protected]> Signed-off-by: ted chang <[email protected]>
foundation-model-stack · Mar 4, 2024 · 896b793 · 896b793
1 parent 87d5fb6
commit 896b793
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/utils/test_data_type_utils.py b/tests/utils/test_data_type_utils.py
@@ -31,4 +31,5 @@ def test_get_torch_dtype():
     for t in dtype_dict.keys():
         # When passed a string, it gets converted to torch.dtype
         assert data_type_utils.get_torch_dtype(t) == dtype_dict.get(t)
+        # When passed a torch.dtype, we get the same torch.dtype returned
         assert data_type_utils.get_torch_dtype(dtype_dict.get(t)) == dtype_dict.get(t)