The GPU memory usage grows with the iteration when calling PyTorch with PyCall, I have opened an issue at github. It seems that PyTorch is unaware of the release of the local variable and I am wondering whether this is related to the Julia’s memory management mechanism. A further simplified code is shown as follows:
Julia
using PyCall
@pyimport torch
device = torch.device(ifelse(torch.cuda[:is_available](), "cuda", "cpu"))
println(device)
for epoch = 1:20
for i = 1:100
Xs = torch.rand(300, 800)[:to](device)
end
sleep(2)
println("Epoch: $epoch")
end
Python
import torch
import time
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
for epoch in range(1, 21):
for i in range(0, 100):
Xs = torch.rand(300, 800).to(device)
time.sleep(2)
print("Epoch: ", epoch)