fix: adapt f16 relu (#76)

ise-uiuc · Feb 14, 2023 · 30743fe · 30743fe
1 parent 23713a3
commit 30743fe
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nnsmith/materialize/torch/numeric.py b/nnsmith/materialize/torch/numeric.py
@@ -12,6 +12,8 @@ def numeric_valid(outputs) -> bool:
 
 # generalized loss fn
 def smoothed_relu(x):
+    if x.dtype == torch.float16:
+        return torch.relu(x.float()).half()
     return torch.relu(x)