Execution without cuda throws error

[bvanaken]

Hi,
when running the script on a machine without cuda support, I'm getting the following error:

File ".../Universal-Transformer-Pytorch/models/UTransformer.py", line 236, in forward
halting_probability = torch.zeros(inputs.shape[0],inputs.shape[1]).cuda()
RuntimeError: torch.cuda.FloatTensor is not enabled.

I suppose the lines 236 - 242 in UTransformer.py require an additional cuda-check.

[andreamad8]

Hi,

you are right, this should be fixed.

Thanks