Training on Custom Dataset using TLT

I have generated KITTI format dataset for 1 class and modified the specs files for single class. Dataset converted successfully into tfrecords but when running the training snippet, i am facing this error. This machine is amzon ec2 instance p2.xlarge with NVIDIA K-80 GPU.

tensorflow.python.framework.errors_impl.InvalidArgumentError: Conv2DCustomBackpropInputOp only supports NHWC

Then i tried on my local machine with nvidia gtx 1060 and training and everything works well.

But I cannot train it locally due to system heat-up. Any help on above error would be appreciated.

The same error and solution in TLT Detectnet TrafficCamNet training not working