If you are using MainModel only for get the mid block it’s not necessary to train it (if i understand the question), so you can do:
opA = optimizer(AuxModel)

with torch.no_grad():
   _, B1 = MainModel(input)
   
opA.zero_grad()
B1 = B1.detach()
prediction= AuxModel(B1)
auxLoss = criterion(prediction)
auxLoss.backward()
opA.step()