(self,
dp_group: dist.ProcessGroup,
num_train_timesteps: int = 1000)
| 70 | class TimestepSamplerMP: |
| 71 | |
| 72 | def __init__(self, |
| 73 | dp_group: dist.ProcessGroup, |
| 74 | num_train_timesteps: int = 1000) -> None: |
| 75 | self.dp_size = dp_group.size() |
| 76 | self.dp_rank = dp_group.rank() |
| 77 | self.train_timesteps = np.arange(num_train_timesteps) |
| 78 | self.t_step_split = np.array_split(self.train_timesteps, self.dp_size) |
| 79 | self.iter = 0 |
| 80 | |
| 81 | def sample_t(self, |
| 82 | batch_size: int, |
nothing calls this directly
no outgoing calls
no test coverage detected