ok it seems there is a problem with the nested_detach function. do your inputs have a labels field (which seems to be the default for self.label_names)?
if that doesn’t solve the problem, my suggestion would be to try debug the prediction_step function by getting a batch of inputs and seeing what happens when you apply nested_detach to them