Here is my code and it caused OOM (“Out Of Memory”) when running:

Python

It firstly caused CUDA to report OOM so I just stupidly removed the “cuda()” to let inference run only on the CPU.

But quickly, the CPU program also reports OOM. And this time I realised that the variable “top10” is an array of tensors, not integers. Therefore I should use “top10[0].item()” to convert it to a pure integer before adding it to the statistics dictionary.

The correct code should be:

Python

Take care of the data type when using PyTorch.