This commit is contained in:
alsaiduq-lab 2025-03-22 18:22:20 -06:00
parent fa22a0d6d7
commit d9d22b7ff7

View file

@ -1,28 +0,0 @@
{ 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
]))
];
};
}