{
  description = "Flake to setup CS2109S";

  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
    flake-utils.url = "github:numtide/flake-utils";
  };

  outputs = { self, nixpkgs, flake-utils }: flake-utils.lib.eachDefaultSystem (system: 
    let pkgs = import nixpkgs {
      inherit system;
      config = {
        allowUnfree = true;
      };
    }; in rec {
      overlays = [(final: prev: {
        pythonPackagesExtensions = [(py-final: py-prev: {
          torch = py-final.torch-bin;
        })];
      })];
      devShell = pkgs.mkShell {
         buildInputs = with pkgs; [
          (python3.withPackages(ps: with ps; 
          let myoptuna = optuna.override { 
            torch = python3.pkgs.torch-bin;
            torchaudio = python3.pkgs.torchaudio-bin;
            torchvision = python3.pkgs.torchvision-bin;
            }; 
          in [
            ipython
            jupyter
            numpy
            pandas
            matplotlib
            torch-bin
            scikit-learn
            timeout-decorator
            torchvision-bin
            seaborn
          ]))
        ];
      };
    }
  );
}