I think this line can be replaced by:
if torch.cuda.is_available():
onehot = torch.cuda.FloatTensor(*x.shape, dims).zero()
else:
onehot = torch.FloatTensor(*x.shape, dims).zero().to(device)
As a result, the cpu usage can be reduced, leading to faster training.
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(): onehot = torch.cuda.FloatTensor(*x.shape, dims).zero() else: onehot = torch.FloatTensor(*x.shape, dims).zero().to(device)
As a result, the cpu usage can be reduced, leading to faster training.