fix

GreenWizard2015 · Apr 12, 2024 · 4c34ea6 · 4c34ea6
1 parent 3a4c3e1
commit 4c34ea6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/NN/utils.py b/NN/utils.py
@@ -34,8 +34,8 @@ def normVec(x):
 
 def ensure4d(src):
   shp = tf.shape(src)
-  if 4 == len(shp): return src
-  if 3 == len(shp): return tf.reshape(src, tf.concat([shp, [1]], axis=0))
+  if 4 == len(src.shape): return src
+  if 3 == len(src.shape): return tf.reshape(src, tf.concat([shp, [1]], axis=0))
   raise ValueError('Invalid shape: {}'.format(shp))
 
 def extractInterpolated(data, pos):