Investigate different sampling behaviour of the norm distribution when the cpu is used as device.

[giorgiapitteri]

When sampling the norm distribution in sample we need to unsqueeze the dimensions when the device is set to gpu. However if the device is set to cpu this unsqueeze needs to be removed otherwise we will get the following error:

r = r.repeat(*[1 for _ in sample_shape], self.dim)
RuntimeError: Number of dimensions of repeat dims can not be smaller than number of dimensions of tensor