maybe good?

cs2109s/labs/flake.nix

@@ -1,5 +1,5 @@
 {
-  description = "A very basic flake";
+  description = "Flake to setup CS2109S";
 
   inputs = {
     nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
@@ -16,23 +16,26 @@
       overlays = [(final: prev: {
         pythonPackagesExtensions = [(py-final: py-prev: {
           torch = py-final.torch-bin;
-          torchvision = py-final.torchvision-bin;
-          torchaudio = py-final.torchaudio-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
+            torch-bin
             scikit-learn
             timeout-decorator
-            torchvision
+            torchvision-bin
             seaborn
           ]))
         ];