loss_hist_cnn = {}

loss_hist_cnn['train'] = {}
loss_hist_cnn['test'] = {}

epochs = 10
for epoch in tqdm(range(epochs)):
    print(f"Epoch: {epoch}n---------")
    loss_cnn = train_step(data_loader=train_loader,
        model=net,
        loss_fn=loss_function,
        optimizer=optimizer,
        accuracy_fn=accuracy_fn
    )

    loss_hist_cnn['train'][epoch] = loss_cnn

    loss_cnn_test = test_step(data_loader=test_loader,
        model=net,
        loss_fn=loss_function,
        accuracy_fn=accuracy_fn
    )

    loss_hist_cnn['test'][epoch] = loss_cnn_test