Thank you for double-checking. The other thread I’d recommend  is

  
    
    
    Torch.cuda.amp cannot speed up on A100 mixed-precision
  
  
    I am trying to use torch.cuda.amp to speed up training. The following code works well on V100 (120s w/o amp and 67s w/ amp), but cannot get a reasonable speedup on A100 (53s w/o amp and 50s w/ amp). 
I am using the most recent NGC docker nvcr.io/nvidia/pytorch:21.04-py3, which has CUDA 11.3, pytorch 1.9.0. 
The strange thing is if I change the model to resnet101, both V100 and A100 can get speedup, but not for the swin transformer on A100. See the following table for the running time (in sec) on…
  


again @ptrblck 's a hero for looking at this.
About the memory: With AMP’s automatic casting so for very small networks, the memory cost of the tensors that are there twice (in FP32 originally and cast to FP16) might be larger than the benefit. (Assuming you use PyTorch’s memory counting, if you use nvidia-smi you’d might also run into caching allocator things.
Best regards
Thomas