Tiny suggestion: make the visualization for torch.zeros and torch.ones have the same y-axis limits so the difference is visually separated.