Utils.checkpoint and cuda.amp, save memory

Yangmin: <p>Hi, I was using cuda.amp.autocast to save memory during training.</p>
<p>But if I use checkpoint in the middle of the network forward pass,<br>
x = checkpoint.checkpoint(self.layer2, x)<br>
feat = checkpoint.checkpoint(self.layer3, x)</p>
<p>the error comes out like below.</p>
<p>RuntimeError: Input type (torch.cuda.HalfTensor) and weight type (torch.cuda.FloatTensor) should be the same</p>
<p>Is it not possible to use both cuda amp and checkpoint?</p>

ptrblck: <p>Yes, it should work, so could you post a minimal, executable code snippet to reproduce the issue, please?</p>

Yangmin: <p>Okay. Below is my model code with layers wrapped with checkpoint.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/3/a3bc559ab08243fa80bd1214fa664099bd646e6f.png" data-download-href="https://discuss.pytorch.org/uploads/default/a3bc559ab08243fa80bd1214fa664099bd646e6f" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/3/a3bc559ab08243fa80bd1214fa664099bd646e6f.png" alt="image" data-base62-sha1="nmtkBbYPV6BEfymrWVgUuzyBIzJ" width="419" height="500" data-dominant-color="100F0D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">566×675 16.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Below is the code for training, wrapped with amp.autocast() and calculating loss with<br>
amp.GradScaler</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e1467f243d879a1fa12e533ade264bbfd74180f4.png" data-download-href="https://discuss.pytorch.org/uploads/default/e1467f243d879a1fa12e533ade264bbfd74180f4" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e1467f243d879a1fa12e533ade264bbfd74180f4.png" alt="image" data-base62-sha1="w8St5GhBcG93wYU42wYf6TFCPje" width="589" height="500" data-dominant-color="100F0E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">718×609 17.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The error comes out in the layer3 of the model where checkpoint is used.<br>
The weight type seems to be not matched when using with amp.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/3/3369385e667ebefa44f7793f2a0b413f5e3a4ac7.png" data-download-href="https://discuss.pytorch.org/uploads/default/3369385e667ebefa44f7793f2a0b413f5e3a4ac7" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/3/3369385e667ebefa44f7793f2a0b413f5e3a4ac7.png" alt="image" data-base62-sha1="7kNMCvuwfltnWN1ab7rLINXJUpN" width="690" height="393" data-dominant-color="1C1C1C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1035×590 31.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Could you post the code snippet by wrapping it into three backticks ```, please?</p>