JIT optimization

[yangsenius]

The NAS search space may contain many for loops in the forward function, which breaks parallel computing, slows down the speed, and restricts the GPU memory utilization.

I have finished for optimizing the class Cell(nn.Module) to Cell(torch.jit.ScriptModule) , and class Meta_Arch(torch.jit.ScriptModule) so on.

Welcome to share your JIT optimization for NAS, particularly for this PoseNFS repo!