| 738 | """ |
| 739 | |
| 740 | def __init__(self, output_size, data_format=None) -> None: |
| 741 | super().__init__() |
| 742 | assert output_size is not None, "'output_size' cannot be NoneType" |
| 743 | self.output_size = _pair(output_size) |
| 744 | if data_format: |
| 745 | if not data_format in ["channels_first", "channels_last"]: |
| 746 | raise ValueError( |
| 747 | f"data_format must be one of ['channels_first', 'channels_last'], but got {data_format}" |
| 748 | ) |
| 749 | self.channel_pos = data_format |
| 750 | elif os.getenv("ONEFLOW_ENABLE_NHWC") == "1": |
| 751 | self.channel_pos = "channels_last" |
| 752 | else: |
| 753 | self.channel_pos = "channels_first" |
| 754 | |
| 755 | def to_memory_format(self, memory_format) -> None: |
| 756 | if memory_format is flow.channels_last: |