search

NeuroDiffGym / neurodiffeq

A library for solving differential equations using neural networks based on PyTorch, used by multiple research groups around the world, including at Harvard IACS.
http://pypi.org/project/neurodiffeq/
MIT License
702 stars 90 forks source link

GPU version? #13

Closed devanshkv closed 4 years ago

devanshkv commented 5 years ago

Hi @feiyu-chen96, do you have plans for the GPU version in the near future?

shuheng-liu commented 4 years ago

Setting the default tensor type right after imports worked for me.

# import libraries
...
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.backends.cudnn.benchmark = True

# normal training
...
devanshkv commented 4 years ago

18 fixes it.