DataParallel fixes (#5733)

* DataParallel fixes:

1. switched to a more precise check
-        if self.args.n_gpu > 1:
+        if isinstance(model, nn.DataParallel):

2. fix tests - require the same fixup under DataParallel as the training module

* another fix
This commit is contained in:
Stas Bekman
2020-07-20 06:29:12 -07:00
committed by GitHub
parent 290b6e18ac
commit 35cb101eae
3 changed files with 7 additions and 2 deletions

View File

@@ -803,6 +803,8 @@ class ModelTesterMixin:
# Wrap model in nn.DataParallel
model = torch.nn.DataParallel(model)
# Our model outputs do not work with DataParallel, so forcing return tuple.
inputs_dict["return_tuple"] = True
with torch.no_grad():
_ = model(**self._prepare_for_class(inputs_dict, model_class))