Solved by my3bikaht in post #4 
              
              
                What I meant is that you don’t need loops to iterate over each dimension. This is very slow indeed and numpy and pytorch have bunch of tricks to evade loops completely, like broadcasting and vectorized operations. 
>>> a = torch.ones(3,3)
>>> a[1,1] = 0
>>> a
tensor([[1., 1., 1.],
        [1., 0., 1…