| 941 | """ |
| 942 | |
| 943 | def __init__(self, output_size, return_indices=False, data_format=None) -> None: |
| 944 | super().__init__(output_size, return_indices=return_indices) |
| 945 | if data_format: |
| 946 | if not data_format in ["channels_first", "channels_last"]: |
| 947 | raise ValueError( |
| 948 | f"data_format must be one of ['channels_first', 'channels_last'], but got {data_format}" |
| 949 | ) |
| 950 | self.channel_pos = data_format |
| 951 | elif os.getenv("ONEFLOW_ENABLE_NHWC") == "1": |
| 952 | self.channel_pos = "channels_last" |
| 953 | else: |
| 954 | self.channel_pos = "channels_first" |
| 955 | |
| 956 | def to_memory_format(self, memory_format) -> None: |
| 957 | if memory_format is flow.channels_last: |