Error in C3W2 zombie detector assignment

Course Q&A TensorFlow: Advanced Techniques Specialization Advanced Computer Vision with TensorFlow
1

I believe that my error is in the train_step_fn, because rest of the things seem aligned

 [CODE REMOVED FROM MENTOR]
the error is

/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/func_graph.py in autograph_handler(*args, **kwargs)
1145 except Exception as e: # pylint:disable=broad-except
1146 if hasattr(e, “ag_error_metadata”):
→ 1147 raise e.ag_error_metadata.to_exception(e)
1148 else:
1149 raise

ValueError: in user code:

File "<ipython-input-55-92fea73694a6>", line 46, in train_step_fn  *
    losses_dict = model.loss(prediction_dict, shapes)
File "/usr/local/lib/python3.10/dist-packages/object_detection/meta_architectures/ssd_meta_arch.py", line 876, in loss  *
    location_losses = self._localization_loss(
File "/usr/local/lib/python3.10/dist-packages/object_detection/core/losses.py", line 78, in __call__  *
    target_tensor = tf.where(tf.is_nan(target_tensor),

ValueError: Shapes must be equal rank, but are 3 and 1 for '{{node Loss/Loss/Select}} = Select[T=DT_FLOAT](Loss/Loss/IsNan, concat_4, Loss/stack_2)' with input shapes: [0], [4,51150,4], [0].
2

Its against conduct rules here to post code solutions! The problem is in the steps before total_loss calculation, there are a bunch of problems with your code.

My suggestions at this point is to re-read again and redo again the exercise fro the beginning.