I agree that the optimization is where the problem is in the code snippet above (even after setting create_graph=True). I looked in to the higher library (https://github.com/facebookresearch/higher 2) which implements a differentiable optimizer. This would allow you to track the updates to the parameters through the many SGD updates. The code would look like this:
with higher.innerloop_ctx(modelB, modelB_optim, copy_initial_weights=False) as (fmodel, diffopt):
    pseudo_labels = modelA(X_train)

    # Train model B on pseudo labels
    logits = fmodel(X_train)
    loss1 = CrossEntropy(logits, pseudo_labels)
    diffopt.step(loss1)

    # Update model A with model B's loss on a validation set
    logits = fmodel(X_val)
    loss2 = CrossEntropy(logits, y_val)
    grad = torch.autograd.grad(loss2, modelA.parameters())

The grad is all zero. I suspect this is because higher takes a snapshot of model B only (to make it functional), the gradients cannot be backprop-ed all the way to model A.
If there was a way to extend the graph for fmodel leading all the way back to model A, this would probably work.