Fix import for newer PyTorch

openclimatefix · Mar 3, 2024 · 4293f2d · 4293f2d
1 parent 7464326
commit 4293f2d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/.ruff.toml b/.ruff.toml
@@ -38,12 +38,15 @@ line-length = 100
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
 # Assume Python 3.10.
-target-version = "py310"
+target-version = "py311"
 fix = false
 # Group violations by containing file.
 format = "github"
 ignore-init-module-imports = true
 
+[format]
+docstring-code-format = true
+
 [mccabe]
 # Unlike Flake8, default to a complexity level of 10.
 max-complexity = 10

diff --git a/environment.yml b/environment.yml
@@ -6,7 +6,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - cudatoolkit=12.1
+  - pytorch-cuda=12.1
   - numcodecs
   - pandas
   - pip