Skip to content

Commit ea36a78

Browse files
authored
Merge pull request #401 from hwangdeyu/deyu/add_HardSwishJitAutoFn_operator
add HardSwishJitAutoFn operator export to onnx
2 parents 4203efa + 7a4be5c commit ea36a78

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

timm/models/layers/activations_me.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,13 @@ def backward(ctx, grad_output):
152152
x = ctx.saved_tensors[0]
153153
return hard_swish_jit_bwd(x, grad_output)
154154

155+
@staticmethod
156+
def symbolic(g, self):
157+
input = g.op("Add", self, g.op('Constant', value_t=torch.tensor(3, dtype=torch.float)))
158+
hardtanh_ = g.op("Clip", input, g.op('Constant', value_t=torch.tensor(0, dtype=torch.float)), g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
159+
hardtanh_ = g.op("Div", hardtanh_, g.op('Constant', value_t=torch.tensor(6, dtype=torch.float)))
160+
return g.op("Mul", self, hardtanh_)
161+
155162

156163
def hard_swish_me(x, inplace=False):
157164
return HardSwishJitAutoFn.apply(x)

0 commit comments

Comments
 (0)