[GPUNet/TorchHub] Updated precision handling

nv-kkudrynski · nv-kkudrynski · commit 9cb7dd07c670 · 2022-05-19T07:33:13.000-07:00
diff --git a/PyTorch/Classification/GPUNet/configs/gpunet_torchhub.py b/PyTorch/Classification/GPUNet/configs/gpunet_torchhub.py
@@ -59,10 +59,6 @@ def nvidia_gpunet(pretrained=True, **kwargs):
 
     modelMath = kwargs.get('model_math', 'fp32')
     if modelMath == "fp16":
-        modelMath = torch.float16
-    else:
-        modelMath = torch.float32
-
-    model = model.to("cuda", modelMath)
+        model.half()
     
     return model
