Add FX based FeatureGraphNet capability

pull/800/head
Alexander Soare 3 years ago
parent 0149ec30d7
commit cf4561ca72

@ -486,4 +486,4 @@ class FeatureGraphNet(nn.Module):
self.graph_module = build_feature_graph_net(model, return_nodes) self.graph_module = build_feature_graph_net(model, return_nodes)
def forward(self, x): def forward(self, x):
return list(self.graph_module(x).values()) return list(self.graph_module(x).values())

Loading…
Cancel
Save