def forward(self, batch):
    df = batch['distance_fields']
    if self.cuda:
        df = df.cuda()
    params = self.model(df[:, None]).view(df.size(0), self.n_primitives, (- 1))
    params = th.cat([(0.35 * params[..., :3]), params[..., 3:]], dim=(- 1))
    df = (utils.distance_to_rounded_cuboids if self.rounded else utils.distance_to_cuboids)
    if self.csg:
        plus_params = params[:, :(self.n_primitives // 2)]
        plus_distance_fields = utils.compute_distance_fields(plus_params, self.canvas_size, df=df)
        plus_distance_fields = plus_distance_fields.min(1)[0]
        minus_params = params[:, (self.n_primitives // 2):]
        minus_distance_fields = utils.compute_distance_fields(minus_params, self.canvas_size, df=df)
        minus_distance_fields = minus_distance_fields.min(1)[0]
        distance_fields = (th.max(plus_distance_fields, (- minus_distance_fields)) ** 2)
    else:
        distance_fields = utils.compute_distance_fields(params, self.canvas_size, df=df)
        distance_fields = (distance_fields.min(1)[0] ** 2)
    alignment_fields = utils.compute_alignment_fields(distance_fields)
    distance_fields = distance_fields[..., 1:(- 1), 1:(- 1), 1:(- 1)]
    occupancy_fields = utils.compute_occupancy_fields(distance_fields)
    return {'distance_fields': distance_fields, 'alignment_fields': alignment_fields, 'occupancy_fields': occupancy_fields}
def forward(self, batch):
    df = batch['distance_fields']
    if self.cuda:
        df = df.cuda()
    params = self.model(df[:, None]).view(df.size(0), self.n_primitives, (- 1))
    params = th.cat([(0.35 * params[..., :3]), params[..., 3:]], dim=(- 1))
    df = (utils.distance_to_rounded_cuboids if self.rounded else utils.distance_to_cuboids)
    if self.csg:
        plus_params = params[:, :(self.n_primitives // 2)]
        plus_distance_fields = utils.compute_distance_fields(plus_params, self.canvas_size, df=df)
        plus_distance_fields = plus_distance_fields.min(1)[0]
        minus_params = params[:, (self.n_primitives // 2):]
        minus_distance_fields = utils.compute_distance_fields(minus_params, self.canvas_size, df=df)
        minus_distance_fields = minus_distance_fields.min(1)[0]
        distance_fields = (th.max(plus_distance_fields, (- minus_distance_fields)) ** 2)
    else:
        distance_fields = utils.compute_distance_fields(params, self.canvas_size, df=df)
        distance_fields = (distance_fields.min(1)[0] ** 2)
    alignment_fields = utils.compute_alignment_fields(distance_fields)
    distance_fields = distance_fields[..., 1:(- 1), 1:(- 1), 1:(- 1)]
    occupancy_fields = utils.compute_occupancy_fields(distance_fields)
    return {'distance_fields': distance_fields, 'alignment_fields': alignment_fields, 'occupancy_fields': occupancy_fields}