Skip to content

GitLab

Commit 25d7d09b authored by mashun1's avatar mashun1
Browse files

Update bias_act.py

parent 451b975a
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
torch_utils/ops/bias_act.py
View file @ 25d7d09b
...@@ -81,8 +81,8 @@ def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None, ...@@ -81,8 +81,8 @@ def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clamp=None,
""" """
assert isinstance(x, torch.Tensor) assert isinstance(x, torch.Tensor)
assert impl in ['ref', 'cuda'] assert impl in ['ref', 'cuda']
if impl == 'cuda' and x.device.type == 'cuda' and _init(): # if impl == 'cuda' and x.device.type == 'cuda' and _init():
return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b) # return _bias_act_cuda(dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp).apply(x, b)
return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp) return _bias_act_ref(x=x, b=b, dim=dim, act=act, alpha=alpha, gain=gain, clamp=clamp)
#---------------------------------------------------------------------------- #----------------------------------------------------------------------------
......
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