Skip to content

Commit

Permalink
Updated logger warnings and function docstring
Browse files Browse the repository at this point in the history
Signed-off-by: Ashvin Kumar <quic_ashvkuma@quicinc.com>
  • Loading branch information
quic-ashvkuma committed Nov 7, 2024
1 parent 1703b12 commit ace95d6
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -746,12 +746,14 @@ def flatten_pack_producers(deconstruct_op: Op):
if len(output_structure) == 1:
self._output_structure = next(iter(output_structure.values()))
else:
logger.warning("Unable to isolate model outputs.")
self._output_structure = None

# Remove inputs called "input_i" and populate their contents into a list based on their index
self._input_structure = [input_structure.pop(f'input_{i}') for i in range(len(input_structure))]
self._input_structure = self._input_structure[0] if len(self._input_structure) == 1 else self._input_structure
if len(input_structure) != 0:
logger.warning("Unable to isolate model inputs.")
self._input_structure = None

@property
Expand Down
6 changes: 4 additions & 2 deletions TrainingExtensions/torch/src/python/aimet_torch/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -663,8 +663,10 @@ def replace_modules_with_instances_of_new_type(model: torch.nn.Module, modules_t
def create_rand_tensors_given_shapes(input_shape: Union[Tuple, List[Tuple]], device: torch.device) \
-> List[torch.Tensor]:
"""
Given shapes of some tensors, create one or more random tensors and return them as a list of tensors
:param input_shape: Shapes of tensors to create
Given shapes of some tensors, create one or more random tensors and return them as a list of tensors. Note that
tensor shapes must expressed as Tuples. A collection of tensor shapes may be represented as List or nested List of
tuples.
:param input_shape: Shapes of tensors to create (expressed as a Tuple, a List of Tuples, or a nested List of Tuples)
:param device: Device to create tensors on
:return: Created list of tensors
"""
Expand Down

0 comments on commit ace95d6

Please sign in to comment.