Add assert for generator layer output shape and defun train_sstep instead of model calls