From 6e631a2da81c304f0be6d817ba0aa446493cddf3 Mon Sep 17 00:00:00 2001 From: Dmytro Doroshenko Date: Sat, 13 Apr 2024 20:00:37 +0300 Subject: [PATCH] fix: move jpeg_quality tensor to device to handle different devices --- kornia/enhance/jpeg.py | 1 + 1 file changed, 1 insertion(+) diff --git a/kornia/enhance/jpeg.py b/kornia/enhance/jpeg.py index d3abca654c..0155f8a1aa 100644 --- a/kornia/enhance/jpeg.py +++ b/kornia/enhance/jpeg.py @@ -526,6 +526,7 @@ def jpeg_codec_differentiable( # Quantization tables to same device and dtype as input image quantization_table_y = quantization_table_y.to(device, dtype) quantization_table_c = quantization_table_c.to(device, dtype) + jpeg_quality = jpeg_quality.to(device) # Perform encoding y_encoded, cb_encoded, cr_encoded = _jpeg_encode( image_rgb=image_rgb,