MyModel(nn.Module):
...
@torch.cuda.amp.autocast()
def forward(self, input):
...
model = MyModel()
dp_model = nn.DataParallel(model)
with torch.cuda.amp.autocast(enable=False): # !!!!!!!!!!!!!
output = dp_model(input)
loss = loss_fn(output)
即使在外面使用了enable=False
, 如果model定义里面,只是在forward上面加@torch.cuda.amp.autocast()
; 还是会启用。
default enable是true。 上层的with torch.cuda.amp.autocast(enable=False):
不能影响到@torch.cuda.amp.autocast()
;
网友评论