Skip to content

Commit

Permalink
allow @conda to accept nvidia as a compute target (#1811)
Browse files Browse the repository at this point in the history
  • Loading branch information
savingoyal authored Apr 24, 2024
1 parent 8b3d566 commit e64c0d2
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion metaflow/plugins/pypi/conda_environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -266,7 +266,9 @@ def get_environment(self, step):
# Resolve `linux-64` Conda environments if @batch or @kubernetes are in play
target_platform = conda_platform()
for decorator in step.decorators:
if decorator.name in ["batch", "kubernetes"]:
# TODO: rather than relying on decorator names, rely on attributes
# to make them extensible.
if decorator.name in ["batch", "kubernetes", "nvidia"]:
# TODO: Support arm architectures
target_platform = "linux-64"
break
Expand Down

0 comments on commit e64c0d2

Please sign in to comment.