Fix lambda function broke torch.save
Summary: Torch.save use pickle, which cannot handle lambda function or local function directly without modify serialization.py. This diff fix the issue by extract lambda to a normal function.
Test Plan: buck test mode/dev-nosan //caffe2/test/fx2trt/core:test_trt_module
Reviewed By: 842974287
Differential Revision: D32320536
fbshipit-source-id: 497d2e64f94526f92e6d1a9909b6ad629dbca850
