Cuda version of torch.eye()

1

What would be the best way to make an identity matrix directly on GPU?
I couldn’t find a torch.cuda.eye() factory function.
I am currently simply making a FloatTensor, then filling it with zeros with .fill_(0) and finally
using torch.masked_scatter_() to put ones on the diagonal.

Any faster/more elegant way to do this?

2

Hi,

The eye function takes a device= argument. You can specify "cuda" there. Or provide a device object.

2 Likes