# Calculate loss function
    loss = criterion(z, data.y)