fix: move jpeg_quality tensor to device to handle different devices

CodSpeedHQ · Apr 13, 2024 · 6e631a2 · 6e631a2
1 parent 2387a2d
commit 6e631a2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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,