if self._sample is not None:
self._sample = self.criterion.prepare(self.model, self._sample)
net_output = self.model(**self._sample["net_input"])
loss_ = self.criterion(net_output, self._sample)
if grad_denom is not None:
loss_ /= grad_denom
loss_.backward()
loss = loss_.data[0]
// flatten grads into a contiguous block of memory
if self.flat_grads is None:
self.flat_grads = self._flatten_grads_(self.model)
After Change
if self._sample is not None:
loss_dict = self.criterion(self.model, self._sample, grad_denom)
loss_dict["loss"].backward()
loss = loss_dict["loss"].data[0]
// flatten grads into a contiguous block of memory
if self.flat_grads is None:
self.flat_grads = self._flatten_grads_(self.model)
