diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 676f63b5c359..02303e261844 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -1509,6 +1509,7 @@ def create_convert_map(self): "softmax": self._softmax, "log_softmax": self._log_softmax, "dropout": lambda node: self.env[node.args[0]], + "stochastic_depth": lambda node: self.env[node.args[0]], "clamp": self._clamp, "relu": lambda node: self.block_builder.emit(relax.op.nn.relu(self.env[node.args[0]])), "leaky_relu": self._leakyrelu, diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index c6c4f2597260..3a5e3509d339 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from torch import fx from torch.nn import Module +import torchvision import tvm from tvm import relax @@ -1160,6 +1161,37 @@ def main( verify_model(Dropout2(), input_info, {}, expected1) +def test_stochastic_depth(): + input_info = [([1, 3, 10, 10], "float32")] + + class StochasticDepth1(Module): + def __init__(self): + super().__init__() + self.stochastic_depth = torchvision.ops.StochasticDepth(0.5, mode="row") + + def forward(self, x): + return self.stochastic_depth(x) + + class StochasticDepth2(Module): + def forward(self, x): + return torchvision.ops.stochastic_depth(x, 0.5, mode="row", training=False) + + @tvm.script.ir_module + class expected1: + @R.function + def main( + input_1: R.Tensor((1, 3, 10, 10), dtype="float32") + ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): + # block 0 + with R.dataflow(): + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = input_1 + R.output(gv) + return gv + + verify_model(StochasticDepth1(), input_info, {}, expected1) + verify_model(StochasticDepth2(), input_info, {}, expected1) + + def test_layernorm(): input_info = [([1, 3, 10, 10], "float32")]