In the step function of GradScaler, if a closure is given as a member of kwargs, RuntimeError occurs by the following codes.

  

      github.com
  

  
    pytorch/pytorch/blob/04fe4382ec8e6031d353530d6b8446ac9c9f507e/torch/cuda/amp/grad_scaler.py#L312-L313 2



      
      
          if "closure" in kwargs:
              raise RuntimeError("Closure use is not currently supported if GradScaler is enabled.")
      
    
  


  

  
    
    
  

  


The error message says “not currently supported”. Are there any plans to support this feature? Or Would you tell me reasons why closure isn’t supported? If possible, I want to try to make a patch to solve this issue.
My motivation is to support SAM with native amp easily.