ERROR: to() received an invalid combination of arguments; training on resnet50 with imagenet dataset

im trying to train a model using resnet 50 on the imagenet dataset but as soon as i run the command provided by repo to train with fp16, i get this error.

$python3 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 ./

`opt_level = O3
keep_batchnorm_fp32 = None <class 'NoneType'>
loss_scale = None <class 'NoneType'>

CUDNN VERSION: 7602

=> creating model 'resnet50'
Traceback (most recent call last):
File "main_amp.py", line 542, in
main()
File "main_amp.py", line 147, in main
model = model.cuda().to(memory_format=memory_format)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 422, in to
device, dtype, non_blocking = torch._C._nn._parse_to(*args, **kwargs)
TypeError: to() received an invalid combination of arguments - got (memory_format=torch.memory_format, ), but expected one of:

• (torch.device device, torch.dtype dtype, bool non_blocking, bool copy)
• (torch.dtype dtype, bool non_blocking, bool copy)
• (Tensor tensor, bool non_blocking, bool copy)`

I currently have PyTorch: 1.2.0, Python 3.6.9. I am also running this on a docker container and ran this command $docker pull pytorch/pytorch:nightly-devel-cuda10.0-cudnn7

Start with main_amp.py at line 147, where the model is moved to CUDA and given a memory format, then reproduce the provided resnet50 FP16 command in the stated Docker environment. Compare the PyTorch 1.2.0 installation with the nightly CUDA 10.0 image and verify that training starts without the reported to() TypeError.