Skip to content

Conversation

@Quentin-Anthony
Copy link

The current imagenet example doesn't move images to the GPU in the distributed GPU case, only in the single GPU case:

https://github.com/pytorch/examples/blob/main/imagenet/main.py#L293
https://github.com/pytorch/examples/blob/main/imagenet/main.py#L337

This leads to errors like the following:

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

I think this can be fixed by relaxing the if statement.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants