run python mnist_visualize.py --model unbounded_stn --angle 90 --grid_size 4 error

[KakaVlasic]

hello! when I run python mnist_visualize.py --model unbounded_stn --angle 90 --grid_size 4, it appears following:

create model with STN
Traceback (most recent call last):
File "mnist_visualize.py", line 47, in
data_list = target2data_list[target]
KeyError: tensor(7)

any help will be appreciate！

[w11m]

I have Same problem

[jackweiwang]

target2data_list = [list() for i in range(10)]

[MingqiangNing]

I think index shouldn't be a tensor, it should be int type, like this :

data_list = target2data_list[int(target)]