How to determine the available GPU memory in Deep Java Library (DJL)?

Question

Is there a way to query the maximum total memory on a GPU in DJL? E.g. an NVIDIA 1080ti should return a value of 11GB.

I am working with the mxnet engine, but an engine-agnostic solution would of course be preferable.

Answer 1

Yes, DJL has integrated with CUDA APIs directly by JNA in https://github.com/awslabs/djl/blob/d2c47d0f2d663b8a5794c21d971420a99f2d47cd/api/src/main/java/ai/djl/util/cuda/CudaUtils.java. So it's engine-agnostic solution.

MemoryUsage mem = CudaUtils.getGpuMemory(device);
mem.getMax(); // it should return 11GB