There are some minor issue. I think this would work for you 
x = torch.zeros(1, 1, 4, 4)
x[:, :, 3] = 1.
x90 = x.transpose(2, 3)
x180 = x.flip(2)
x270 = x.transpose(2, 3).flip(3)
print(x)
print(x90)
print(x180)
print(x270)
2 Likes