Cool! Your fix works. You saved my day (almost)!
I’ve encountered “out of memory” crash using caffe to extract features with pretrained resnet. So I rewrote the code in pytorch and still met this error. After wrapping the extraction part in a function, it went much further: from 400+ images to 1200+ images. However, in the end it still ran out of memory.
I was watching the GPU RAM usage closely, and saw the RAM increased from 4G to 9G and then 11G before it crashed (the total GPU RAM is 12G).
Really confused why this would happen. Anyway I’m just using some standard model to do minimal processing.
Update:
I managed to reduce RAM use by wrapping the extraction code in a no_grad block:
with torch.no_grad():
    features = model(im).data

Now the code runs through all the 15000+ images without any error 
Though, I still don’t understand why keeping the gradients will cause the memory problem, as I’m looping through each individual image, and in the next loop, the gradients should be dereferenced and released automatically.
Update 2:
Finally I solved the memory problem! I realized that in each iteration I put the input data in a new tensor, and pytorch generates a new computation graph. That causes the used RAM to grow forever. Then I use a placeholder tensor and copy the data to this tensor, and the RAM always stays at a low level