How to put flow to gpu?

gitlabspy commented 4 years ago

I use realnvp flow in examples folder and I tinkered it a little bit. It seems that it only works on cpu.

class FlowModel(Flow):
    def __init__(...):
        transformations = CompositeTransform([...])
        super().__init__(
            transform=CompositeTransform([...]),
            distribution=StandardNormal((3, 32, 32)),
        )

flow = FlowModel()
flow = flow.to(device) # CUDA:0 here
...
flow.transform_to_noise(z) # error...

error:

File "/data/lhb/anaconda2/envs/python36/lib/python3.6/site-packages/nflows/transforms/base.py", line 52, in _cascade
    total_logabsdet += logabsdet
RuntimeError: expected backend CUDA and dtype Float but got backend CPU and dtype Float

I am pretty new to pytorch but have some experience in tensorflow, so how to put flow to gpu to train? Thx!!

andersdot commented 4 years ago

I was having a similar issue, I tried

torch.set_default_tensor_type('torch.cuda.FloatTensor')

which the developers recommended in another issue. They mention it's a hack, but it worked for me. Cheers

andersdot commented 4 years ago

Not a great hack, as they suggested, I'm having OOM issues.

arturbekasov commented 4 years ago

Thanks for reporting this! #9 has been merged now, which should hopefully have solved most of GPU-related issues. Feel free to re-open if your particular issue hasn't been resolved by this, letting us know what distribution/transform you're using exactly. Cheers.

bayesiains / nflows

How to put flow to gpu? #15