{ pkgs ? import (fetchTarball https://nixos.org/channels/nixos-unstable/nixexprs.tar.xz) {
  config = {
    allowUnfree = true;
    cudaSupport = true;
  };
} }:

with pkgs;

mkShell {
  buildInputs = [
    (python3.withPackages(p: [p.pytorch-bin p.numpy]))
  ];
}