These images are the images showing model evaluation and inference. When the step svi.run is encountered my CPU memory shoots up to 100 percent and in google colab the CUDA is out of memory. I reduced the batch size to as small as one image per batch. But no help.
Also I tried running torch.cuda.empty_cache(). It did not help either