From 17f15a6ce4a3edc9bd65d90672ae0d518120fb9e Mon Sep 17 00:00:00 2001
From: Alessandro Pappalardo <1934033+volcacius@users.noreply.github.com>
Date: Thu, 13 Jul 2023 23:00:21 +0200
Subject: [PATCH] Fix (QuantTensor): pixel_shuffle and unshuffle handler (#663)

* Fix (QuantTensor): pixel_shuffle and unshuffle handler

* fixup
---
 src/brevitas/quant_tensor/torch_handler.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/brevitas/quant_tensor/torch_handler.py b/src/brevitas/quant_tensor/torch_handler.py
index 5708fd35f..6b530a320 100644
--- 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)
\ No newline at end of file