Thanks for your reply.
As suggested, the train works well without memory issue.
But, in load_params(avg_param) at the test session, I got the following error.
param.data.copy_(flattened[offset:offset + param.nelement()]).view(param.size())
RuntimeError: copy from list to FloatTensor isn't implemented

I think the load_param function should be modified due to the list.