| 226 | """ |
| 227 | |
| 228 | def __init__( |
| 229 | self, |
| 230 | kernel_size: _size_2_t, |
| 231 | stride: Optional[_size_2_t] = None, |
| 232 | padding: _size_2_t = 0, |
| 233 | dilation: _size_2_t = 1, |
| 234 | return_indices: bool = False, |
| 235 | ceil_mode: bool = False, |
| 236 | ): |
| 237 | super().__init__() |
| 238 | self.kernel_size = _pair(kernel_size) |
| 239 | self.stride = _pair(stride) if (stride is not None) else _pair(kernel_size) |
| 240 | self.padding = _pair(padding) |
| 241 | self.dilation = _pair(dilation) |
| 242 | self.return_indices = return_indices |
| 243 | self.ceil_mode = ceil_mode |
| 244 | if os.getenv("ONEFLOW_ENABLE_NHWC") == "1": |
| 245 | self.channel_pos = "channels_last" |
| 246 | else: |
| 247 | self.channel_pos = "channels_first" |
| 248 | |
| 249 | def to_memory_format(self, memory_format) -> None: |
| 250 | if memory_format is flow.channels_last: |