28 lines
578 B
Nix
28 lines
578 B
Nix
{ config, pkgs, lib, ... }:
|
|
|
|
with lib;
|
|
|
|
let
|
|
cfg = config.custom.torch;
|
|
in {
|
|
options.custom.torch = {
|
|
enable = mkEnableOption "PyTorch with CUDA support";
|
|
cudaSupport = mkOption {
|
|
type = types.bool;
|
|
default = true;
|
|
};
|
|
package = mkOption {
|
|
type = types.package;
|
|
default = pkgs.python311Packages.torch.override { cudaSupport = cfg.cudaSupport; };
|
|
};
|
|
};
|
|
|
|
config = mkIf cfg.enable {
|
|
environment.systemPackages = with pkgs; [
|
|
(python3.withPackages (ps: with ps; [
|
|
cfg.package
|
|
numpy
|
|
]))
|
|
];
|
|
};
|
|
}
|