Fix (QuantTensor): pixel_shuffle and unshuffle handler (#663)

* Fix (QuantTensor): pixel_shuffle and unshuffle handler * fixup
Xilinx · Jul 13, 2023 · 17f15a6 · 17f15a6
1 parent b8a90af
commit 17f15a6
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/brevitas/quant_tensor/torch_handler.py b/src/brevitas/quant_tensor/torch_handler.py
@@ -153,4 +153,9 @@ def interpolate_handler(
 
 @implements(F.pixel_shuffle)
 def pixel_shuffle_handler(*args, **kwargs):
-    return quant_invariant_handler(F.pixel_shuffle_handler, *args, **kwargs)
+    return quant_invariant_handler(F.pixel_shuffle, *args, **kwargs)
+
+
+@implements(F.pixel_unshuffle)
+def pixel_unshuffle_handler(*args, **kwargs):
+    return quant_invariant_handler(F.pixel_unshuffle, *args, **kwargs)