{
  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs?ref=23.05-pre";
    fenix.url = "github:nix-community/fenix";
  };

  outputs = { self, nixpkgs, flake-utils, fenix }:
    flake-utils.lib.eachDefaultSystem (system:
      let
        pkgs = import nixpkgs {
          inherit system;
          config.allowUnfree = true;
        };

        torch = pkgs.python310Packages.torch.override { cudaSupport = true; };

        torchvision =
          pkgs.python310Packages.torchvision.override { inherit torch; };
      in rec {
        devShell = pkgs.mkShell {
          packages = (with pkgs.python310Packages; [ ])
            ++ [ torch torchvision ];
        };
      });
}