(self, timestep)
| 227 | |
| 228 | |
| 229 | def training_weight(self, timestep): |
| 230 | |
| 231 | timestep_id = torch.argmin((self.timesteps - timestep |
| 232 | .unsqueeze(-1).repeat(1, self.timesteps.shape[0]) |
| 233 | .to(self.timesteps.device)).abs(), dim=1) |
| 234 | # timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) |
| 235 | weights = self.linear_timesteps_weights[timestep_id] |
| 236 | return weights.to(timestep.device) |