From 7c5d7f14e1e5948078dfe017b3e097c0071974f6 Mon Sep 17 00:00:00 2001
From: billera <lukasbillera01@gmail.com>
Date: Mon, 29 Apr 2024 18:24:50 +0200
Subject: [PATCH] reshape T

---
 src/layers.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/layers.jl b/src/layers.jl
index 34997c1..585b345 100644
--- a/src/layers.jl
+++ b/src/layers.jl
@@ -247,7 +247,7 @@ function ipa_customgrad(ipa::Union{IPCrossA, IPA}, Ti::Tuple{AbstractArray,Abstr
     else
         use_softmax1 = false
     end
-    TiL = TiR = Ti 
+    TiL = TiR = (Ti[1], reshape(Ti[2], size(Ti[2],1), 1, size(Ti[2])[2:end]...)) 
     siL = siR = S
     rot_TiL, translate_TiL = TiL
     rot_TiR, translate_TiR = TiR