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 · 87d5fb6 · 87d5fb6
1 parent 863cc53
commit 87d5fb6
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
@@ -29,5 +29,6 @@ def test_str_to_torch_dtype_exit():
 
 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)
         assert data_type_utils.get_torch_dtype(dtype_dict.get(t)) == dtype_dict.get(t)