A tensor is called i-type if its entries are equal to the sum of index squared.
I want to have a function to construct an i-type tensor with arbitrarily given shape. This can be done with numpy as follows:
import numpy as np
def i_tensor(shape = [2,2,2]):
    t = np.zeros(shape)
    it = np.nditer(t, flags=['multi_index'])
    while not it.finished:
        ind = it.multi_index
        t[ind] = np.sum(np.array(ind)**2)
        it.iternext()
    return t


Test 1:
#test the function
> i_tensor([4,4])

> array([[ 0.,  1.,  4.,  9.],
        [ 1.,  2.,  5., 10.],
        [ 4.,  5.,  8., 13.],
        [ 9., 10., 13., 18.]])

Test2:
> i_tensor([2,2,3])
> array([[[0., 1., 4.],
        [1., 2., 5.]],

       [[1., 2., 5.],
        [2., 3., 6.]]])

Can I have the same function for torch tensor? I do not see tensor replacement of nditer.