is there a way to programmatically clear CUDA context?
I would like to release the memory allocated by CUDA context after model training because my process need to do other operations after that

  
    
    
    Model.to("cpu") does not release GPU memory allocated by registered buffer 
  
  
    I have figured that registered_buffer does not release GPU memory when the model is moved back to CPU. 
Here is the minimal code for reproducing the observation 
import torch
from torch import nn
from subprocess import Popen, PIPE

class TestNet(nn.Module):
    def __init__(
        self
    ):
        super().__init__()
        self.register_buffer("test", torch.Tensor([1,2,3]).view(-1, 1, 1), False)

def log_nvidia_smi():
    print("--- nvidia-smi ---")
    for line in Popen('nvidia-smi', shel…