Getting the deep learning environment with all NVidia libaries set up is a bit fiddly. Thankfully I can lean on Nix for that as well. I have a package definition like this (dl.nix):

{ python, pythonPackages, gcc, cudatoolkit75, cudnn5_cudatoolkit75 }:

python.buildEnv.override {
  extraLibs = with pythonPackages; [
    (Keras.override { propagatedBuildInputs = [ Theano-cuda six pyyaml future ];  })
  extraPaths = [ gcc cudatoolkit75 cudnn5_cudatoolkit75 ];
  ignoreCollisions = true;

(Note that this needs this branch in your nixpkgs. I’m hoping it’ll be merged soon.)

With that in place it’s trivial to drop into a nix-shell that then has a python to run your nix code:

$ nix-shell -p "callPackage ./dl.nix {}"

I’m also still on Theano rather than Tensorflow because I find it easier to debug. Tensorflow available in nixpkgs though if you prefer that.