Skip to content

GitLab

Commit 8b6f8fc1 authored by Thor Johnsen's avatar Thor Johnsen
Browse files

Bug fix

parent 64b93e3e
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
apex/contrib/bottleneck/bottleneck.py
View file @ 8b6f8fc1
...@@ -508,7 +508,7 @@ class SpatialBottleneckFunction(torch.autograd.Function): ...@@ -508,7 +508,7 @@ class SpatialBottleneckFunction(torch.autograd.Function):
btm_relu_halo = relu1[:,:,Hs-1:,:].clone() btm_relu_halo = relu1[:,:,Hs-1:,:].clone()
btm_grad_out1 = grad_out1[:,:,Hs-1:,:] btm_grad_out1 = grad_out1[:,:,Hs-1:,:]
w1by3 = w[:,:1,:,:].clone() w1by3 = w[:,:1,:,:].clone()
ctx.stream1.wait_stream(ctx.stream2) # wait for halo transfers to finish ctx.stream2.wait_stream(ctx.stream1) # wait for halo transfers to finish
ctx.stream2.wait_stream(torch.cuda.current_stream()) # wait for backward_grad_out1_mask to finish before launching halo correction kernel ctx.stream2.wait_stream(torch.cuda.current_stream()) # wait for backward_grad_out1_mask to finish before launching halo correction kernel
with torch.cuda.stream(ctx.stream1): with torch.cuda.stream(ctx.stream1):
btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(ctx.explicit_nhwc, ctx.stride_1x1, t_list, w1by3, grads, btm_halo, btm_relu_halo, btm_grad_out1.clone()) btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo_corr(ctx.explicit_nhwc, ctx.stride_1x1, t_list, w1by3, grads, btm_halo, btm_relu_halo, btm_grad_out1.clone())
......
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