quick fix #5978

[jackalcooper]

No description provided.

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

PyTorch resnet50 time: 140.1ms (= 7003.6ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
OneFlow resnet50 time: 128.1ms (= 6406.0ms / 50, input_shape=[16, 3, 224, 224], backward is enabled)
Relative speed: 1.09 (= 140.1ms / 128.1ms)

PyTorch resnet50 time: 84.9ms (= 4246.6ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
OneFlow resnet50 time: 74.8ms (= 3741.5ms / 50, input_shape=[8, 3, 224, 224], backward is enabled)
Relative speed: 1.13 (= 84.9ms / 74.8ms)

PyTorch resnet50 time: 57.1ms (= 2854.4ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
OneFlow resnet50 time: 47.3ms (= 2364.1ms / 50, input_shape=[4, 3, 224, 224], backward is enabled)
Relative speed: 1.21 (= 57.1ms / 47.3ms)

PyTorch resnet50 time: 47.0ms (= 2350.8ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
OneFlow resnet50 time: 39.6ms (= 1979.4ms / 50, input_shape=[2, 3, 224, 224], backward is enabled)
Relative speed: 1.19 (= 47.0ms / 39.6ms)

PyTorch resnet50 time: 37.2ms (= 1861.3ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
OneFlow resnet50 time: 37.3ms (= 1865.4ms / 50, input_shape=[1, 3, 224, 224], backward is enabled)
Relative speed: 1.00 (= 37.2ms / 37.3ms)