Getting error in Trainer function for action recognition

epochs = 65
lr_decay_count = 0
valid_loss_best = 1000
valid_acc_best = 10
hist_prec_train =
hist_prec_test =
hist_recall_train =
hist_recall_test =
hist_loss =
hist_loss_valid =

for epoch in range(epochs):
tic = time.time()
train_metric.reset()
valid_metric.reset()
train_loss = 0
valid_loss = 0

# Learning rate decay
if epoch == lr_decay_epoch[lr_decay_count]:
    trainer.set_learning_rate(trainer.learning_rate*lr_decay)
    lr_decay_count += 1

# Loop through each batch of training data
y_true = np.array([], dtype='int')
y_pred = np.array([], dtype='int')
for i, batch in enumerate(train_data):
    # Extract data and label
    data = mx.gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0,even_split=False)
    label =mx.gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0,even_split=False)

    #AutoGrad
    with ag.record():
       output = []
       for _, X in enumerate(data):
          X = X.reshape((-1,) + X.shape[2:])
          pred = net(X)
          output.append(pred)

    loss = [loss_fn(yhat, y) for yhat, y in zip(output, label)]

    # Backpropagation
    for l in loss:
        l.backward()

    # Optimize
    trainer.step(batch_size)

    # Update metrics
    train_loss += sum([l.mean().asscalar() for l in loss])
    train_metric.update(label, output)

    y_true = np.concatenate((y_true, label[0].asnumpy()))
    y_pred = np.concatenate((y_pred, pred.argmax(axis=1).astype('int').asnumpy()))

name, acc = train_metric.get()
precisions = metrics.precision_score(y_true, y_pred, average=None, zero_division=False)
recall = metrics.recall_score(y_true, y_pred, average=None, zero_division=False)

# Update history and print metrics
train_history.update([acc])
print(f'[Epoch {epoch}] train={acc:.4f} loss={train_loss/(i+1):.4f} time: {time.time()-tic:.1f} sec')
print('Train precision: ',{k:v for k,v in zip(classes, precisions)})
print('Train recall: ',{k:v for k,v in zip(classes, recall)})
hist_loss.append(train_loss/(i+1))
hist_prec_train.append({k:v for k,v in zip(classes, precisions)})
hist_recall_train.append({k:v for k,v in zip(classes, recall)})

y_true_v = np.array([], dtype='int')
y_pred_v = np.array([], dtype='int')
for i, batch in enumerate(valid_data):
    # Extract data and label
    data = mx.gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0,even_split=False)
    label = mx.gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0,even_split=False)

    output = []
    for _, X in enumerate(data):
        X = X.reshape((-1,) + X.shape[2:])
        pred = net(X)
        output.append(pred)
    loss = [loss_fn(yhat, y) for yhat, y in zip(output, label)]

    # Update metrics
    valid_loss += sum([l.mean().asscalar() for l in loss])
    valid_metric.update(label, output)
    y_true_v = np.concatenate((y_true_v, label[0].asnumpy()))
    y_pred_v = np.concatenate((y_pred_v, pred.argmax(axis=1).astype('int').asnumpy()))

name, acc = valid_metric.get()
precisions_v = metrics.precision_score(y_true_v, y_pred_v, average=None, zero_division=False)
recall_v = metrics.recall_score(y_true_v, y_pred_v, average=None, zero_division=False)

# Update history and print metrics
valid_history.update([acc])
print(f'valid_acc: {acc}, valid_loss: {valid_loss/(i+1)}')
hist_loss_valid.append(valid_loss/(i+1))
print(f'valid precision:', {k:v for k,v in zip(classes, precisions_v)})
print(f'valid recall:', {k:v for k,v in zip(classes, recall_v)})
hist_prec_test.append({k:v for k,v in zip(classes, precisions_v)})
hist_recall_test.append({k:v for k,v in zip(classes, recall_v)})

if (valid_loss_best > valid_loss) or (valid_acc_best < acc) :
    valid_loss_best = valid_loss
    valid_acc_best = acc
    print(f'Best valid loss: {valid_loss_best}')
    file_name = f"{models_path}/i3d_resnet50_v1_custom{epoch}.params"
    net.save_parameters(file_name)

ValueError Traceback (most recent call last)
in ()
38 output.append(pred)
39
—> 40 loss = [loss_fn(yhat, y) for yhat, y in zip(output, label)]
41
42 # Backpropagation

5 frames
/usr/local/lib/python3.7/dist-packages/mxnet/ndarray/ndarray.py in reshape(self, *shape, **kwargs)
1522 # Array size should not change
1523 if np.prod(res.shape) != np.prod(self.shape):
→ 1524 raise ValueError(‘Cannot reshape array of size {} into shape {}’.format(np.prod(self.shape), shape))
1525 return res
1526

ValueError: Cannot reshape array of size 5 into shape (1, 2)