flake/modules/aipython3/default.nix
2023-02-12 01:04:57 +01:00

29 lines
601 B
Nix

{ lib, ... }:
{
perSystem = { pkgs, ... }: {
dependencySets = let
overlays = import ./overlays.nix pkgs;
mkPythonPackages = overlayList: let
python3' = pkgs.python3.override {
packageOverrides = lib.composeManyExtensions overlayList;
};
in python3'.pkgs;
in {
aipython3-amd = mkPythonPackages [
overlays.fixPackages
overlays.extraDeps
overlays.torchRocm
];
aipython3-nvidia = mkPythonPackages [
overlays.fixPackages
overlays.extraDeps
overlays.torchCuda
];
};
};
}