Your loss calculation is wrapped into a no_grad block:
with torch.no_grad():
    XX = torch.mean(kernels[:batch_size, :batch_size])
    YY = torch.mean(kernels[batch_size:, batch_size:])
    XY = torch.mean(kernels[:batch_size, batch_size:])
    YX = torch.mean(kernels[batch_size:, :batch_size])
    loss = torch.mean(XX + YY - XY - YX)

will will make sure that Autograd isn’t tracking these operations.
Could you disable the nn.CrossEntropyLoss, only use the mmd_loss, call backward on it, and check the model parameters for valid gradients?