pipeline last stage supports multi output (#151)

2025-09-02 17:46:42 +00:00 · 2022-01-17 15:57:47 +08:00
parent 1ff5be36c2
commit 7bf1e98b97
2 changed files with 35 additions and 18 deletions
--- a/tests/test_trainer/test_pipeline/model/resnet.py
+++ b/tests/test_trainer/test_pipeline/model/resnet.py
@@ -139,7 +139,7 @@ class VanillaResNet(ModelFromConfig):
    def forward(self, x: Tensor):
        for layer in self.layers:
            x = layer(x)
-        return x,
+        return x

    def init_weights(self):
        for m in self.modules():