You can do this with Cog! Once you've written cog.yaml, you can run arbitrary commands inside the environment which has CUDA installed correctly:

  $ cog run python train.py