Hello,
Here’s the current code where I turn on QAT only near the end of training, and still has the same issue:
import torch
from torch import nn, optim
from torch.quantization import QuantStub, DeQuantStub
from copy import deepcopy

print(torch.__version__)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.backbone0 = nn.Sequential(
            QuantStub(),
            nn.Conv2d(3, 1, 1, bias=False),
            nn.BatchNorm2d(1),
            nn.ReLU(),
            DeQuantStub(),
        )
        self.backbone1 = nn.Sequential(
            nn.Conv2d(1, 2, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(2),
            nn.MaxPool2d(14),
            nn.Sigmoid(),
        )

    def forward(self, x):
        x = self.backbone0(x)
        x = self.backbone1(x)
        return x

model = Model()

# torch.quantization.fuse_modules(model, [['1', '2', '3'], ['4', '5']], inplace=True)

# model.backbone0.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
# torch.quantization.prepare_qat(model.backbone0, inplace=True)

optimizer = optim.Adam(model.parameters(), lr=1)
model.to(device)

criterion = nn.BCELoss()
for epoch in range(1000):
    # print('EPOCH', epoch)
    model.train()

    inputs = torch.rand(2, 3, 28, 28)
    labels = torch.FloatTensor([[1, 1], [0, 0]])

    inputs = inputs.to(device)
    labels = labels.to(device)
    outputs = model(inputs)
    loss = criterion(outputs.view(2, 2), labels)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if epoch == 945: # turn on qat
        model.to('cpu')
        model.backbone0.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
        torch.quantization.prepare_qat(model.backbone0, inplace=True)
        model.to(device)

    if epoch >= 950:
        model.apply(torch.quantization.disable_observer)
        pass

    if epoch >= 950:
        model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)

    if epoch == 999:
        # print('MODEL', model)
        quant_model = deepcopy(model)
        quant_model = torch.quantization.convert(quant_model.eval().cpu(), inplace=False)

        with torch.no_grad():
            inp = torch.zeros([1, 3, 28, 28], device='cpu')
            model.eval().cpu()
            quant_model.eval().cpu()

            qout = quant_model.forward(inp)
            out = model.forward(inp)

            print(f"<============== EPOCH {epoch} ===============>")
            print(out.view(2).tolist(), "#not converted")
            print(qout.view(2).tolist(), "#quant converted")
            print(f"<============== NOT CONVERTED MODEL ===============>")
            print(model.state_dict())
            print(f"<============== QUANT CONVERTED MODEL ===============>")
            print(quant_model.state_dict())

            model.to(device)