{ config, lib, pkgs, ...}: {
  # Make sure opengl is enabled
  hardware.opengl = {
    enable = true;
    driSupport = true;
    driSupport32Bit = true;
  };
  
  # Tell Xorg to use the nvidia driver (also valid for Wayland)
  services.xserver.videoDrivers = [
    # For the RTX 4000
    "nvidia"
  ];

  services.xserver.exportConfiguration = true;
  hardware.nvidia = {
    
    # Modesetting is needed for most Wayland compositors
    modesetting.enable = true;

    # Use the open source version of the kernel module
    # Only available on driver 515.43.04+
    # Seems to cause some issues with Tensorflow/CUDA/PyTorch
    open = false;

    # Enable the nvidia settings menu
    nvidiaSettings = true;

    # Optionally, you may need to select the appropriate driver version for your specific GPU.
    package = config.boot.kernelPackages.nvidiaPackages.stable;
  };

  # Allow Unfree, Nvidia, CUDA, and derived packages.
  nixpkgs.config.allowUnfree = true;
  nixpkgs.config.cudaSupport = true;

  # Add cuda to system environment
  environment.systemPackages = with pkgs; [
    cudatoolkit
  ];

  # Use Prebuilt Cached Versions of CUDA Packages
  nix.settings = {
    substituters = [
      "https://cuda-maintainers.cachix.org"
    ];
    trusted-public-keys = [
      "cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
    ];
  };
}