Fix optim load_state_dict bug

[wyg1997]

No description provided.

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

OneFlow resnet50 time: 128.7ms (= 6434.0ms / 50, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 139.8ms (= 6991.0ms / 50, input_shape=[16, 3, 224, 224])
Relative speed: 1.09 (= 139.8ms / 128.7ms)

OneFlow resnet50 time: 74.8ms (= 3740.6ms / 50, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 84.3ms (= 4213.8ms / 50, input_shape=[8, 3, 224, 224])
Relative speed: 1.13 (= 84.3ms / 74.8ms)

OneFlow resnet50 time: 49.1ms (= 2455.4ms / 50, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 58.1ms (= 2904.5ms / 50, input_shape=[4, 3, 224, 224])
Relative speed: 1.18 (= 58.1ms / 49.1ms)

OneFlow resnet50 time: 40.5ms (= 2023.9ms / 50, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 49.0ms (= 2450.2ms / 50, input_shape=[2, 3, 224, 224])
Relative speed: 1.21 (= 49.0ms / 40.5ms)

OneFlow resnet50 time: 46.8ms (= 2341.5ms / 50, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 44.7ms (= 2232.9ms / 50, input_shape=[1, 3, 224, 224])
Relative speed: 0.95 (= 44.7ms / 46.8ms)

OneFlow resnet50 time: 143.9ms (= 7194.9ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 162.3ms (= 8115.3ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
Relative speed: 1.13 (= 162.3ms / 143.9ms)

OneFlow resnet50 time: 93.8ms (= 4692.1ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 103.7ms (= 5185.8ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
Relative speed: 1.11 (= 103.7ms / 93.8ms)

OneFlow resnet50 time: 68.0ms (= 3399.1ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 73.0ms (= 3652.5ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
Relative speed: 1.07 (= 73.0ms / 68.0ms)

OneFlow resnet50 time: 64.6ms (= 3227.9ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 69.3ms (= 3465.3ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
Relative speed: 1.07 (= 69.3ms / 64.6ms)

OneFlow resnet50 time: 67.7ms (= 3385.7ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 58.9ms (= 2942.5ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
Relative speed: 0.87 (= 58.9ms / 67.7ms)