Fixing NaN warning with single element tensors
@@ -93,6 +93,7 @@ class TensorboardWriter(FromParams):
# Log parameter values to Tensorboard
for name, param in model.named_parameters():
self.add_train_scalar("parameter_mean/" + name, param.data.mean())
+ if param.data.numel() > 1:
self.add_train_scalar("parameter_std/" + name, param.data.std())
if param.grad is not None:
if param.grad.is_sparse:
@@ -104,6 +105,7 @@ class TensorboardWriter(FromParams):
# skip empty gradients
if torch.prod(torch.tensor(grad_data.shape)).item() > 0: # pylint: disable=not-callable
self.add_train_scalar("gradient_mean/" + name, grad_data.mean())
+ if grad_data.numel() > 1:
self.add_train_scalar("gradient_std/" + name, grad_data.std())
else:
# no gradient for a parameter with sparse gradients
