I have a code that looks like this:
    for epoch in range(epochs):
        aloss_t,aloss_E_t,aloss_F_t,MAE_t,Fps_t,Fts_t, t_dataload_t, t_prepare_t, t_model_t, t_backprop_t = use_model(model, dataloader_train, train=True, max_samples=1e6, optimizer=optimizer, device=device, batch_size=batch_size)
        aloss_v,aloss_E_v,aloss_F_v,MAE_v,Fps_v,Fts_v,t_dataload_v, t_prepare_v, t_model_v, t_backprop_v = use_model(model, dataloader_val, train=False, max_samples=100, optimizer=optimizer, device=device, batch_size=batch_size)

where use_model is a function that can either be used for training or validation over a dataset. The model runs fine during training and uses about 70% of my gpu memory with batchsize=40, however when it exists the training loop and switches to validation the memory does not get freed, and hence it runs out of memory when running the validation.
I thought the memory would be automatically freed when I left the training function, but apparently not, how do I ensure that this happens?