Replaced torch.randn(1)[1] by torch.normal() call