Speedup for training

[SongtaoLiu0823]

https://github.com/molecule-one/megan/blob/master/src/model/megan.py#L19

I think this line can be replaced by:
if torch.cuda.is_available():
one_hot = torch.cuda.FloatTensor(*x.shape, dims).zero_()
else:
one_hot = torch.FloatTensor(*x.shape, dims).zero_().to(device)

As a result, the cpu usage can be reduced, leading to faster training.

[SongtaoLiu0823]

In my machine, the CPU usage is from 2000+ to 100+.

[mikolajsacha]

Hi, sorry for the late response, and thank you for your find. We will merge it into the codebase :)