I tried your code and you are right. It does takes around 2 seconds if you measure it using the timeit.Timer function. But if you try it with just PyTorch code as in:
import torch
a = torch.zeros(10, 10, 20)
Then it’s pretty much instantaneous. Seems like the timeit.Timer function adds some weird time overhead in this case. Maybe someone with more knowledge than me could clarify why that is the case.
Best regards,
diego