Out of memory error: Multi-GPUs not used when training semi-supervised model

[W-XuWX]

Hi I was trying to follow the "unsupervised_hello_world" notebook tutorial, but trained on Stanford cars dataset. (https://github.com/tensorflow/similarity/blob/master/examples/unsupervised_hello_world.ipynb).

Running into out of memory error even with 4x 16gb Nvidia V100 as only one gpu is utilised. (Using AWS ec2 instance)

---

Error message:

---

nvidia-smi:

---

Relevant code snippet: