staged

modules/default.nix

@@ -19,6 +19,7 @@
     ./steam.nix
     ./timezone.nix
     ./user.nix
+    ./torch.nix
   ];
 
   # Enable CUPS to print documents.

modules/torch.nix

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