Fix extra gpu memory in flow.load and module.load_state_dict

[daquexian]

No description provided.

[github-actions]

Code got formatted by CI. Please request CI again if you still want to have this PR merged. If the PR is from a forked repo, please download the patch files from the GitHub Actions web page and apply them locally.

[github-actions]

View latest API docs preview at: https://staging.oneflow.info/docs/Oneflow-Inc/oneflow/pr/8301/

[github-actions]

Speed stats:

GPU Name: NVIDIA GeForce GTX 1080 

❌ OneFlow resnet50 time: 130.4ms (= 13036.6ms / 100, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 143.4ms (= 14335.3ms / 100, input_shape=[16, 3, 224, 224])
✔️ Relative speed: 1.10 (= 143.4ms / 130.4ms)

OneFlow resnet50 time: 76.6ms (= 7655.5ms / 100, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 85.4ms (= 8540.1ms / 100, input_shape=[8, 3, 224, 224])
✔️ Relative speed: 1.12 (= 85.4ms / 76.6ms)

OneFlow resnet50 time: 53.9ms (= 10772.7ms / 200, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 62.9ms (= 12586.8ms / 200, input_shape=[4, 3, 224, 224])
✔️ Relative speed: 1.17 (= 62.9ms / 53.9ms)

OneFlow resnet50 time: 43.2ms (= 8638.6ms / 200, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 41.2ms (= 8232.2ms / 200, input_shape=[2, 3, 224, 224])
❌ Relative speed: 0.95 (= 41.2ms / 43.2ms)

OneFlow resnet50 time: 35.8ms (= 7159.4ms / 200, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 42.1ms (= 8422.0ms / 200, input_shape=[1, 3, 224, 224])
✔️ Relative speed: 1.18 (= 42.1ms / 35.8ms)

OneFlow swin dataloader time: 0.251s (= 50.169s / 200, num_workers=1)
PyTorch swin dataloader time: 0.153s (= 30.697s / 200, num_workers=1)
Relative speed: 0.612 (= 0.153s / 0.251s)

OneFlow swin dataloader time: 0.065s (= 13.007s / 200, num_workers=4)
PyTorch swin dataloader time: 0.041s (= 8.265s / 200, num_workers=4)
Relative speed: 0.635 (= 0.041s / 0.065s)

OneFlow swin dataloader time: 0.036s (= 7.271s / 200, num_workers=8)
PyTorch swin dataloader time: 0.022s (= 4.469s / 200, num_workers=8)
Relative speed: 0.615 (= 0.022s / 0.036s)

❌ OneFlow resnet50 time: 146.4ms (= 14640.1ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 176.3ms (= 17627.8ms / 100, input_shape=[16, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.20 (= 176.3ms / 146.4ms)

OneFlow resnet50 time: 96.4ms (= 9644.8ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 112.8ms (= 11280.9ms / 100, input_shape=[8, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.17 (= 112.8ms / 96.4ms)

OneFlow resnet50 time: 74.8ms (= 14956.1ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 88.4ms (= 17679.4ms / 200, input_shape=[4, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.18 (= 88.4ms / 74.8ms)

OneFlow resnet50 time: 61.2ms (= 12240.2ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 74.2ms (= 14846.9ms / 200, input_shape=[2, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.21 (= 74.2ms / 61.2ms)

OneFlow resnet50 time: 54.4ms (= 10879.1ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 69.0ms (= 13792.0ms / 200, input_shape=[1, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.27 (= 69.0ms / 54.4ms)