//  [python]
def _trainStep(self, datasetIdx, layerIdx):
        #loss function & regularization
        data = tf.switch_case(datasetIdx, self.trainDatasetSelector)

        inputTransformSelector = {}
        for x in range(self.nInputLayers):
            inputTransformSelector[x] = lambda x=x: self.inputLayers[x](
                data['inputFeatures'], training=True)

        regLossSelector = {}
        for x in range(self.nInputLayers):
            regLossSelector[x] = lambda x=x: self.inputLayers[x].losses

        with tf.GradientTape() as tape:
            inputTransformedFeatures = tf.switch_case(
                layerIdx, inputTransformSelector)
            predictions = self.model(inputTransformedFeatures, training=True)
            regularization_loss = tf.math.add_n(self.model.losses) + \
                tf.math.add_n(tf.switch_case(layerIdx, regLossSelector))

            batchSize = tf.shape(data['inputFeatures'])[0]
            if self.args['lossType'] == 'ctc':
                sparseLabels = tf.cast(tf.sparse.from_dense(
                    data['seqClassIDs']), dtype=tf.int32)
                sparseLabels = tf.sparse.SparseTensor(
                    indices=sparseLabels.indices,
                    values=sparseLabels.values-1,
                    dense_shape=[batchSize, self.args['dataset']['maxSeqElements']])

                nTimeSteps = self.model.getSubsampledTimeSteps(data['nTimeSteps'])
                pred_loss = tf.compat.v1.nn.ctc_loss_v2(sparseLabels,
                                                        predictions,
                                                        None,
                                                        nTimeSteps,
                                                        logits_time_major=False,
                                                        unique=None,
                                                        blank_index=-1,
                                                        name=None)

                pred_loss = tf.reduce_mean(pred_loss)

            elif self.args['lossType'] == 'ce':
                mask = tf.tile(data['ceMask'][:, :, tf.newaxis], [
                                1, 1, self.args['dataset']['nClasses']])
                ceLoss = tf.keras.losses.CategoricalCrossentropy(
                    from_logits=True)
                pred_loss = ceLoss(
                    data['classLabelsOneHot'][:,:,0:-1], predictions[:, :, 0:-1]*mask)
                pred_loss = tf.reduce_mean(pred_loss)
                newClassSignalError = tf.reduce_mean(tf.math.square(tf.math.sigmoid(predictions[:,:,-1]) - data['newClassSignal']))
                pred_loss += newClassSignalError

            total_loss = pred_loss + regularization_loss

        #compute gradients + clip
        grads = tape.gradient(total_loss, self.trainableVariables)
        grads, gradNorm = tf.clip_by_global_norm(
            grads, self.args['gradClipValue'])

        #only apply if gradients are finite and we are in train mode
        allIsFinite = []
        for g in grads:
            if g != None:
                allIsFinite.append(tf.reduce_all(tf.math.is_finite(g)))
        gradIsFinite = tf.reduce_all(tf.stack(allIsFinite))

        if gradIsFinite:
            self.optimizer.apply_gradients(zip(grads, self.trainableVariables))

        #compute sequence-element error rate (edit distance) if we are in validation & ctc mode
        #return interval activations so we can visualize what's going on
        #intermediate_output = self.model.getIntermediateLayerOutput(inputTransformedFeatures)

        output = {}
        output['logits'] = predictions
        #output['rnnUnits'] = intermediate_output
        output['inputFeatures'] = data['inputFeatures']
        if self.args['lossType'] == 'ce':
            output['classLabels'] = data['classLabelsOneHot']
        output['predictionLoss'] = pred_loss
        output['regularizationLoss'] = regularization_loss
        output['gradNorm'] = gradNorm
        output['seqIDs'] = data['seqClassIDs']
        output['seqErrorRate'] = tf.constant(0.0)

        return output