self.stats = {key: (val.mean(), val.std()) for key, val in self.data.items() 
                      if type(val) is torch.Tensor and val.dim() == 1 and val.is_floating_point()}