Can not reproduce stably from ckpt

[ChrisDong-THU]

When I was using the SubMConv2d module, I found that the ckpt saved during training could not reproduce its performance after loading it during evaluation. Here is my testing code:

if __name__ == "__main__":
    bs = 4
    num_gs = 256
    in_channels = 64
    out_channels = 64
    kernel_size = 5
    fm_shape = (32, 32)
    use_out_proj = False
    
    load = True
    
    import pytorch_lightning as pl
    pl.seed_everything(0)
    torch.set_float32_matmul_precision("high")
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    spconv2d = SparseConv(in_channels, out_channels, kernel_size, fm_shape, use_out_proj)
    
    if load:
        sd = torch.load('./tmp/spconv2d.pth')
        spconv2d.load_state_dict(sd)
    else:
        torch.save(spconv2d.state_dict(), './tmp/spconv2d.pth')
    
    spconv2d = spconv2d.cuda()
    spconv2d = spconv2d.eval()
    
    instance_feature = torch.randn(bs, num_gs, in_channels).cuda()
    anchor = torch.randn(bs, num_gs, 5 + in_channels).cuda()
    
    with torch.no_grad():
        output = spconv2d(instance_feature, anchor)

SubMConv2d is a layer of SparseConv, whose forward func I've carefully checked for debugging:

input = SparseConvTensor(instance_feature.flatten(0, 1), indices=batch_indices, spatial_shape=spatial_shape, batch_size=b)
output = self.conv(input)

The input stayed the same, but output varied between two possible outcomes. You can easily see this by running the above demo. Looking forward to a solution.