Fair enough - but could we at least get rid of the need for X.cpu().numpy()? Seems X.numpy() alone should be enough.
X.cpu().numpy()
X.numpy()