Have problems in sampling image

[Kaiwen Hong]

Good repo! When I try to run training_transformer.py and want to generate image using line 59, I ran in to problem: there is no codebook in self.vqgan.codebook in transformer.py so I assume I should comment this line (198) and uncomment line 199.
However, if I do this, I will run into problems:
, thread: [30,0,0] Assertion srcIndex < srcSelectDimSize failed.
../aten/src/ATen/native/cuda/Indexing.cu:1141: indexSelectLargeIndex: block: [164,0,0], thread: [31,0,0] Assertion srcIndex < srcSelectDimSize failed...
Would you mind clarifying what should I do for sampling images? I use default hyperparameters and duplicate one image from landscape dataset for 12 times as dataset.

and here's the error log.
Traceback (most recent call last):
File "training_transformer.py", line 143, in
train_transformer = TrainTransformer(args)
File "training_transformer.py", line 36, in init
self.train(args)
File "training_transformer.py", line 48, in train
logits, target = self.model(imgs)
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/kaiwen/Desktop/interview/copilot4d/test/transformer.py", line 57, in forward
_, z_indices = self.encode_to_z(x)
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
return func(*args, **kwargs)
File "/home/kaiwen/Desktop/interview/copilot4d/test/transformer.py", line 42, in encode_to_z
quant_z, , (, _, indices) = self.vqgan.encode(x)
File "/home/kaiwen/Desktop/interview/copilot4d/test/vq_f16.py", line 27, in encode
h = self.encoder(x)
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/kaiwen/Desktop/interview/copilot4d/test/vq_modules.py", line 267, in forward
hs = [self.conv_in(x)]
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 463, in forward
return self._conv_forward(input, self.weight, self.bias)
File "/home/kaiwen/anaconda3/envs/ddpm/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 459, in _conv_forward
return F.conv2d(input, weight, bias, self.stride,
RuntimeError: cuDNN error: CUDNN_STATUS_MAPPING_ERROR