Method __init__

python/oneflow/test/graph/test_graph_pipeline.py:201–216

(self, data)

def _train_with_module(iter_num=3, data=None):
    class DataModule(flow.nn.Module):
        def __init__(self, data):
            super().__init__()
            self.data_list = []
            self.idx = 0
            for pair in data:
                for i in range(4):
                    s = i * 4
                    e = s + 4
                    micro_batch_image = pair[0][s:e]
                    micro_batch_label = pair[1][s:e]
                    self.data_list.append(
                        (
                            flow.Tensor(micro_batch_image).to("cuda:3"),
                            flow.Tensor(micro_batch_label).to("cuda:3"),
                        )
                    )

        def forward(self):
            image = self.data_list[self.idx][0]