[Bug]: - All VRAM of first device gets allocated when importing the lib on a gpu environment.

[Foglia-m]

Module

None

Current Behavior

When running import xplique on a gpu enabled environment, the vram of the first device is allocated at maximum (32go / 32go)

All modules from the init are concerned except from . import plots

Expected Behavior

I would expect lazy allocatino of the ram only when I load models. (is it due to the recent addition of the torch wrapper?)

Version

Version: 1.3.0

Environment

- OS: linux 
- Python version: 3.9.16
- Tensorflow version: 2.8.2

Relevant log output

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.60.13    Driver Version: 525.60.13    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100S-PCI...  On   | 00000000:00:07.0 Off |                    0 |
| N/A   33C    P0    38W / 250W |  31452MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100S-PCI...  On   | 00000000:00:08.0 Off |                    0 |
| N/A   33C    P0    41W / 250W |    474MiB / 32768MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

To Reproduce

import xplique in a python console / notebook

[AntoninPoche]

Hello @Foglia-m, thank you for your contribution. We are going to investigate this ASAP, but a quick fix may be:

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

import xplique

Can you try this and tell us if it helps!

Can you also tell if you had this issue on previous Xplique versions?