dazzle-me:

Is there any way we can prevent CPU bottlenecking the GPU in case of small models?


Assuming the data loading is the bottleneck, you could take a look at this post 10, which explains common use cases and potential workarounds.
If not already done, use multiple workers in the DataLoader.