Skip to content

GitLab

Unverified Commit 1a43f292 authored by Thor Johnsen's avatar Thor Johnsen Committed by GitHub
Browse files

Merge pull request #1329 from NVIDIA/leave_bottleneck_masks_as_bool 

Leave bottleneck masks as bool
parents 7e1c22d0 bd7c1a0f
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
apex/contrib/bottleneck/bottleneck.py
View file @ 1a43f292
...@@ -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