Skip to content

Cog installs wrong pytorch version when cuda is specified #436

Description

@dashstander

OS: Ubuntu 20.04

GPU: RTX 3090

cog version: 0.0.20

With this minimal cog.yaml I run into problems.

build: 
    gpu: true
    cuda: "11.3"
    python_version: "3.8"
    python_packages:
        - "torch==1.10.2"
        - "torchvision==0.11.3"

The image builds without error, just a warning about being unsure whether cuda and pytorch match. However, this matches the pytorch recommended installation versions as of 2022/02/22 which is to run pip3 install torch==1.10.2+cu113 torchvision==0.11.3+cu113 torchaudio==0.10.2+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html.

However, once the image is built I cannot actually use the GPU, the issue being that running:

cog run python -c "import torch; print(torch.__version__)"

gives me "1.10.2+cu102", and cuda 10.2 is not compatible with the architecture of the 3090.

To get around this I added the command pytorch recommends above directly to the run section of the cog.yaml, which works but is very clunky.

Metadata

Metadata

Assignees

No one assigned

    Labels

    type/bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions