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

for img_batch in data_load:
	input = torch.from_numpy(np.asarray(img_batch)).cuda()
    result = self._net(input.permute(0, 3, 1, 2).float())
    values, indices = torch.topk(result, 10)
    for index in range(len(values)):
    	top10 = values[0]
    	statistics["accumulate"] += top10[0]

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:

...
    for index in range(len(values)):
    	top10 = values[0]
    	statistics["accumulate"] += top10[0].item()

Take care of the data type when using PyTorch.