sagemaker-debugger icon indicating copy to clipboard operation
sagemaker-debugger copied to clipboard

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

Open Cpruce opened this issue 6 years ago • 0 comments

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

Cpruce avatar Dec 14 '19 01:12 Cpruce