CouplingFlow not working on GPU

[vpratz]

For some reason, the variables of the RandomPermutation are not moved to GPU. This leads to the following error, where TensorFlow tries to access variables from the CPU when compiling for the GPU:

2024-12-09 14:27:40.186935: W tensorflow[/core/framework/op_kernel.cc:1841](http://127.0.0.1:8899/core/framework/op_kernel.cc#line=1840)] OP_REQUIRES failed at xla_ops.cc:515 :+1:
INVALID_ARGUMENT: Trying to access resource coupling_flow[/random_permutation/variable/43](http://127.0.0.1:8899/random_permutation/variable/43) [...] located in device [/job](http://127.0.0.1:8899/job):localhost[/replica:0](http://127.0.0.1:8899/replica#line=-1)[/task:0](http://127.0.0.1:8899/task#line=-1)[/device](http://127.0.0.1:8899/device):CPU:0 from device [/job](http://127.0.0.1:8899/job):localhost[/replica:0](http://127.0.0.1:8899/replica#line=-1)[/task:0](http://127.0.0.1:8899/task#line=-1)[/device](http://127.0.0.1:8899/device):GPU:0

To reproduce, run the TwoMoons_StarterNotebook.ipynb on a GPU with the TensorFlow backend, and replace the inference network with a standard coupling flow: inference_network = bf.networks.CouplingFlow().

I suppose we somewhere interrupt the tracking of which variables belong to the network, maybe it is in the construction of the CouplingFlow.