Skip to content

install torch and subdependencies at runtime pending compute type - #51

Draft
chelseatroy wants to merge 2 commits into
mainfrom
ctroy-fix-img-classifier
Draft

install torch and subdependencies at runtime pending compute type#51
chelseatroy wants to merge 2 commits into
mainfrom
ctroy-fix-img-classifier

Conversation

@chelseatroy

Copy link
Copy Markdown
Collaborator

Goal(s) of this Pull Request:

Example of how it works:

...

Acceptance criteria:

In order to approve, reviewer must be able to...

  • ...

@chelseatroy
chelseatroy marked this pull request as draft October 22, 2025 18:45
import torch
assert not torch.cuda.is_available(), "This step should be CPU-only."

import torchvision

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@chelseatroy I believe the solution here is not to hack around the torchvision-cpu install, but rather move the whole transform dance to the GPU-enable steps or to use the NVIDIA decorator here too.

For non-trivial uses, we'd still want to do dataset processing on a GPU.

@pypi(
python="3.11.9",
packages={"torch": "2.4.1", "torchvision": "0.19.1", "mozmlops": "0.1.4"},
packages={"mozmlops": "0.1.4"} # intentionally omit torch/torchvision here

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need to omit the dependencies here?

RuntimeImportHelper.ensure_cpu_torch()

import torch
assert not torch.cuda.is_available(), "This step should be CPU-only."

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@chelseatroy the 'nvct' decorator will be discontinued after November. If moving all the steps to nvct solves the problem, let's do that instead

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants