From 2220d4ed2f70ec77027a02f8221f0c7591c4de7c Mon Sep 17 00:00:00 2001 From: Andrew Millspaugh Date: Wed, 21 Apr 2021 12:34:47 -0700 Subject: [PATCH] Support wrapping and optional lowering on model Summary: Productionization/testing requires model output in dict. This change supports exporting the model as such. Differential Revision: D27917487 fbshipit-source-id: a651b47d5398ce52e19446cda2d7417cd7ca1274 --- classy_vision/models/classy_block.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/classy_vision/models/classy_block.py b/classy_vision/models/classy_block.py index cff9b1776a..363a0b6e4c 100644 --- a/classy_vision/models/classy_block.py +++ b/classy_vision/models/classy_block.py @@ -17,13 +17,10 @@ class ClassyBlock(nn.Module): def __init__(self, name, module): super().__init__() self.name = name - self.output = torch.zeros(0) self._module = module def wrapped_module(self): return self._module def forward(self, input): - output = self._module(input) - self.output = output - return output + return self._module(input)