Update src/model_utils.py

src/model_utils.py

@@ -11,7 +11,7 @@ from parallelism_utils import estimate_zero1_model_states_mem_needs, estimate_ze
 
 DTYPE_MODIFIER = {"float32": 1, "float16/bfloat16": 2, "int8": 4, "int4": 8}
 PRECISION = {"Mixed precision": "mixed", "Single precision": "single"}
-DTYPE = {"float32": torch.float32, "float16/bfloat16": torch.float16}
+DTYPE = {"float32": "float32", "float16/bfloat16": "float16"}
 
 
 def extract_from_url(name: str):