diff --git a/pkgs/development/python-modules/augmax/default.nix b/pkgs/development/python-modules/augmax/default.nix
index 879bffb8a0d0..3be4c4713600 100644
--- a/pkgs/development/python-modules/augmax/default.nix
+++ b/pkgs/development/python-modules/augmax/default.nix
@@ -1,6 +1,6 @@
 { buildPythonPackage
 , einops
-, fetchPypi
+, fetchFromGitHub
 , jax
 , jaxlib
 , lib
@@ -15,19 +15,21 @@ buildPythonPackage rec {
 
   disbaled = pythonOlder "3.6";
 
-  # Using fetchPypi because the latest version was not tagged on GitHub.
-  # Switch back to fetchFromGitHub when a tag will be available
-  # https://github.com/khdlr/augmax/issues/8
-  src = fetchPypi {
-    inherit pname version;
-    hash = "sha256-pf1DTaHA7D+s2rqwwGYlJrJOI7fok+WOvOCtZhOOGHo=";
+  src = fetchFromGitHub {
+    owner = "khdlr";
+    repo = "augmax";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-xz6yJiVZUkRcRa2rKZdytfpP+XCk/QI4xtKlNaS9FYo=";
   };
 
-  nativeBuildInputs = [
+  build-system = [
     setuptools
   ];
 
-  propagatedBuildInputs = [ einops jax ];
+  dependencies = [
+    einops
+    jax
+  ];
 
   # augmax does not have any tests at the time of writing (2022-02-19), but
   # jaxlib is necessary for the pythonImportsCheckPhase.