Update handler.py

handler.py

@@ -8,18 +8,21 @@ import time
 import os
 import torch
 
-def max_pooling(model_output):    
-    # Get dimensions    
-    _, Z, Y = model_output.shape    
-    # Initialize an empty list with length Y (384 in your case)    
-    output_array = [0] * Y    
-    # Loop over secondary arrays (Z)    
-    for i in range(Z):        
-        # Loop over values in innermost arrays (Y)        
-        for j in range(Y):            
-            # If value is greater than current max, update max            
-            if model_output[0][i][j] > output_array[j]:                
-                output_array[j] = model_output[0][i][j]    
+def max_pooling(model_output):
+    # Get dimensions
+    Z, Y = len(model_output[0]), len(model_output[0][0])
+  
+    # Initialize an empty list with length Y (384 in your case)
+    output_array = [0] * Y 
+
+    # Loop over secondary arrays (Z)
+    for i in range(Z): 
+        # Loop over values in innermost arrays (Y)
+        for j in range(Y): 
+            # If value is greater than current max, update max
+            if model_output[0][i][j] > output_array[j]: 
+                output_array[j] = model_output[0][i][j]
+    
     return output_array
     
 #Mean Pooling - Take attention mask into account for correct averaging