Skip to content

GitLab

Commit bd7c1a0f authored by Thor Johnsen's avatar Thor Johnsen
Browse files

Leave bottleneck masks as bool

parent 7e1c22d0
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
apex/contrib/bottleneck/bottleneck.py
View file @ bd7c1a0f
...@@ -36,14 +36,14 @@ class FrozenBatchNorm2d(torch.nn.Module): ...@@ -36,14 +36,14 @@ class FrozenBatchNorm2d(torch.nn.Module):
@torch.jit.script @torch.jit.script
def drelu_dscale1(grad_o, output, scale1): def drelu_dscale1(grad_o, output, scale1):
relu_mask = (output>0).half() relu_mask = (output>0)
dx_relu = relu_mask * grad_o dx_relu = relu_mask * grad_o
g1 = dx_relu * scale1 g1 = dx_relu * scale1
return g1, dx_relu return g1, dx_relu
@torch.jit.script @torch.jit.script
def drelu_dscale2(grad_o, output, scale1, scale2): def drelu_dscale2(grad_o, output, scale1, scale2):
relu_mask = (output>0).half() relu_mask = (output>0)
dx_relu = relu_mask * grad_o dx_relu = relu_mask * grad_o
g1 = dx_relu * scale1 g1 = dx_relu * scale1
g2 = dx_relu * scale2 g2 = dx_relu * scale2
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment