From 7fc87def1ea90164967d0d9e22c7de5c7d3c190c Mon Sep 17 00:00:00 2001
From: ChUrl <christoph.urlacher@protonmail.com>
Date: Sun, 13 Nov 2022 14:26:04 +0100
Subject: [PATCH] add torch rocm support to env

---
 flake.nix | 36 +++++++++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/flake.nix b/flake.nix
index 5a9c3e2..2bd15a2 100644
--- a/flake.nix
+++ b/flake.nix
@@ -14,15 +14,45 @@
           overlays = [ devshell.overlay ];
         };
 
+        # TODO: Originally it was nixpkgs.fetchurl but that didn't work, pkgs.fetchurl did...
+        #       Determine the difference between nixpkgs and pkgs
+
+        # Taken from: https://github.com/gbtb/nix-stable-diffusion/blob/master/flake.nix
+        # Overlay: https://nixos.wiki/wiki/Overlays
+        # FetchURL: https://ryantm.github.io/nixpkgs/builders/fetchers/
+        torch-rocm = pkgs.hiPrio (pkgs.python310Packages.torch-bin.overrideAttrs (old: {
+          src = pkgs.fetchurl {
+            name = "torch-1.12.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
+            url = "https://download.pytorch.org/whl/rocm5.1.1/torch-1.12.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
+            hash = "sha256-kNShDx88BZjRQhWgnsaJAT8hXnStVMU1ugPNMEJcgnA=";
+          };
+        }));
+        torchvision-rocm = pkgs.hiPrio (pkgs.python310Packages.torchvision-bin.overrideAttrs (old: {
+          src = pkgs.fetchurl {
+            name = "torchvision-0.13.1+rocm5.1.1-cp310-cp310-linux_x86_64.whl";
+            url = "https://download.pytorch.org/whl/rocm5.1.1/torchvision-0.13.1%2Brocm5.1.1-cp310-cp310-linux_x86_64.whl";
+            hash = "sha256-mYk4+XNXU6rjpgWfKUDq+5fH/HNPQ5wkEtAgJUDN/Jg=";
+          };
+        }));
+
         myPython = pkgs.python310.withPackages (p: with p; [
+          # Basic
+          rich
+
+          # Discord
           discordpy
           python-dotenv
+          pynacl
+
+          # Scraping
           beautifulsoup4
           requests
-          pynacl
-          rich
-          torch
+
+          # MachineLearning
+          torch-rocm
+          torchvision-rocm
           numpy
+          matplotlib
           nltk
         ]);
       in {