Move self._log_params(module) from forward_pre_hook to forward_hook for Pytorch

Not all parameters have been created until after the first step if creating parameters via tracing (during runtime). Can confirm this works thanks to Rahul H:

def forward_hook(self, module, inputs, outputs):
if not self._get_collections_to_save_for_step():
return

self._log_params(module)
module_name = self.module_maps[module]
# This overwhelms the logs; turn back on if you really need it
# logger.debug("Processing the global step {0} for module {1}".format(self.step, module_name))

# Output input tensor
self._write_inputs(module_name, inputs)

# Output output tensors
self._write_outputs(module_name, outputs)
self.last_saved_step = self.step