In my case, my layerA is much more complicated than a nn.Linear layer, so simply detaching weightings would not be practical. I am thinking of deepcopying layerA instead, however im not familiar with the interactions of it with pytorch. Would you know if this will work?
from copy import deepcopy
import torch

x = torch.randn(1, 10)
layerA = ...
layerB = ...
layerA_2 = deepcopy(layerA)

out = layerA(x)
out = layerB(out)
out = layerA_2(out)

out.mean().backward()