Run with GPU

[llbbcc]

We can run the code on the GPU with the following modifications:

pip install onnxruntime-gpu

Modify the following code segment in ./imwatermark/rivaGAN.py

RivaWatermark.encoder = onnxruntime.InferenceSession(
    os.path.join(modelDir, 'rivagan_encoder.onnx'))
RivaWatermark.decoder = onnxruntime.InferenceSession(
    os.path.join(modelDir, 'rivagan_decoder.onnx'))

to

RivaWatermark.encoder = onnxruntime.InferenceSession(
    os.path.join(modelDir, 'rivagan_encoder.onnx'), providers=['CUDAExecutionProvider'])
RivaWatermark.decoder = onnxruntime.InferenceSession(
    os.path.join(modelDir, 'rivagan_decoder.onnx'), providers=['CUDAExecutionProvider'])

Then, you can run the code on the GPU.