Update activations.py

This minor change makes swish and mish inplace by default and saves up to 20% of memory in my tests.
pull/156/head
bonlime 5 years ago committed by GitHub
parent 5966654052
commit 62fea452a3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -18,7 +18,7 @@ if _USE_MEM_EFFICIENT_ISH:
# recomputing torch.sigmoid(x) in backward instead of saving it.
@torch.jit.script
def swish_jit_fwd(x):
return x.mul(torch.sigmoid(x))
return x.mul_(torch.sigmoid(x))
@torch.jit.script
@ -50,7 +50,7 @@ if _USE_MEM_EFFICIENT_ISH:
@torch.jit.script
def mish_jit_fwd(x):
return x.mul(torch.tanh(F.softplus(x)))
return x.mul_(torch.tanh(F.softplus(x)))
@torch.jit.script

Loading…
Cancel
Save