Training the model

There are a number of steps to be performed in training that require a custom training loop. First, let's define a method that executes one step of the training loop. This method is defined in the s2s-training.py file:

@tf.function
def train_step(inp, targ, enc_hidden, max_gradient_norm=5):
    loss = 0
    
    with tf.GradientTape() as tape:
        # print("inside gradient tape")
        enc_output, enc_hidden = encoder(inp, enc_hidden)
               
        dec_hidden = enc_hidden
        dec_input = tf.expand_dims([start] * BATCH_SIZE, 1)
        
        # Teacher forcing - feeding the target as the next input
        for t in range(1, targ.shape[1]):
            # passing enc_output to the decoder
            predictions, dec_hidden, _ = decoder(dec_input,   
                                           dec_hidden, enc_output)
            
            loss += s2s.loss_function(targ[:, t], predictions)
            # using teacher forcing...