How to set batch size in Google Colab notebook according to allocated node specs

It is related to this question. @Vladimir Sotnikov 's answer gave insctruction to get the config information of the allocated colab node. He also showed how to get such information with pytorch and suggested a separate question for it. And here you go.

Solution

it is from Vladimir Sotnikov's answer:

import torch
batch_size=int(torch.cuda.get_device_properties(0).total_memory/1e9+.5)*16
print(f'GPU mem:{int(torch.cuda.get_device_properties(0).total_memory/1e9+.5)}G, batch_size:{batch_size}')