26 lines
680 B
Nix
26 lines
680 B
Nix
|
{
|
||
|
|
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 ];
|
||
|
|
};
|
||
|
|
});
|
||
|
|
}
|