何时在pytorch中分配GPU内存？

Question

例如，我们可以有如下代码：

model.train()
...
optimizer.zero_grad()
...
trn_x, trn_y = next(train_loader)
val_x, val_y = next(valid_loader)
trn_x, trn_y = trn_x.cuda(), trn_y.cuda()
...
logits_search_train, _ = model_small(trn_x)
loss_weight = criterion(logits_search_train, trn_y)
loss_weight.backward()
optimizer.step()

此代码中的哪里有GPU内存分配/取消分配？