Add FX based FeatureGraphNet capability

3 years ago · cf4561ca72
parent 0149ec30d7
commit cf4561ca72
1 changed files with 1 additions and 1 deletions
--- a/timm/models/fx_features.py
+++ b/timm/models/fx_features.py
@ -486,4 +486,4 @@ class FeatureGraphNet(nn.Module):
        self.graph_module = build_feature_graph_net(model, return_nodes)
    def forward(self, x):
-        return list(self.graph_module(x).values())
+        return list(self.graph_module(x).values())