@ptrblck, thanks for your reply.  I’m using amp.
Here is a code snippet to reproduce the issue.
import torch
import torch.nn as nn
from apex import amp
import torch.distributed as dist


class SomeModel(nn.Module):
    def __init__(self):
        super(SomeModel, self).__init__()
        self.conv = nn.Conv3d(
            3,
            16,
            kernel_size=(1, 3, 3),
            stride=1,
            padding=(0, 1, 1),
            bias=False)
        self.bn1 = nn.BatchNorm3d(16)
        self.relu = nn.ReLU(inplace=True)
        self.avgpool = nn.AdaptiveAvgPool3d(1)
        self.fc = nn.Linear(16, 3)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x



print('init process group')
dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:7001',
                            world_size=1, rank=0)

model = SomeModel().cuda()
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, )
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
print('ddp')
model = torch.nn.parallel.DistributedDataParallel(model, find_unused_parameters=True)

print('model train')
# model.train() # works

model.eval()
model.module.fc.train()

x = torch.randn((5, 3, 7, 7, 7), device='cuda')
y = torch.ones((5, ), device='cuda').long()

print('model forward')
outputs = model(x)
print('calculate loss')
loss = criterion(outputs, y)
print('model backward')
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
print('optimizer step')
optimizer.step()

Also, while I make the code snippet, I found that BN cause the issue.
Without BN, no error raised, though I’m not sure it works properly as intended.