//  [python]
def train(self):
        self.writer = SummaryWriter(log_dir=os.path.join(self.base_exp_dir, 'logs'))
        self.update_learning_rate()
        res_step = self.end_iter - self.iter_step

        for iter_i in tqdm(range(res_step)):
            self.hair_network.embed_fn_fine.current_iter = self.iter_step
            self.sdf_network.embed_fn_fine.current_iter = self.iter_step
            self.color_network.embedview_fn.current_iter = self.iter_step
            self.nerf_outside.embed_fn.current_iter = self.iter_step
            self.nerf_outside.embed_fn_view.current_iter = self.iter_step

            data, cam_intr_0, cam_pose_0, image_index = self.sample_rays()

            pixels, radii, true_rgb, mask, hair_mask, orient_angle, orient_conf = data.split([3, 1, 3, 1, 1, 1, 1], dim=1)
            
            if self.camera_model is not None and self.iter_step > self.tune_cameras_start:
                cam_intr_, cam_pose_ = self.camera_model(image_index, cam_intr_0, cam_pose_0)
            else:
                cam_intr_ = cam_intr_0
                cam_pose_ = cam_pose_0

            p = torch.matmul(torch.linalg.inv(cam_intr_)[None, :3, :3], pixels[:, :, None]).squeeze() if len(cam_intr_.shape) < 3 else torch.matmul(torch.linalg.inv(cam_intr_)[:, :3, :3], pixels[:, :, None]).squeeze() # batch_size, 3
            rays_d = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
            rays_d = torch.matmul(cam_pose_[None, :3, :3], rays_d[:, :, None]).squeeze() if len(cam_pose_.shape) < 3 else torch.matmul(cam_pose_[:, :3, :3], rays_d[:, :, None]).squeeze()# batch_size, 3
            if len(cam_pose_.shape) < 3: 
                rays_o = cam_pose_[None, :3, 3].expand(rays_d.shape) # batch_size, 3
            elif cam_pose_.shape[0] == 1:
                rays_o = cam_pose_[:, :3, 3].expand(rays_d.shape) 
            else:
                rays_o = cam_pose_[:, :3, 3]

            cam_intr = cam_intr_.detach()
            cam_extr = torch.linalg.inv(cam_pose_.detach())


            near, far = self.dataset.near_far_from_sphere(rays_o.detach(), rays_d.detach())

            
            background_rgb = None
            if self.use_white_bkgd:
                background_rgb = torch.ones([1, 3])

            if self.mask_weight > 0.0:
                if self.conf['train']['binarize_gt_masks']:
                    mask = (mask > 0.5).float()
            else:
                mask = torch.ones_like(mask)
            
            if  self.hair_mask_weight > 0.0:
                if self.conf['train']['binarize_gt_masks']:
                    hair_mask = (hair_mask > 0.5).float()
            else:
                hair_mask = torch.ones_like(mask)

            mask_sum = mask.sum() + 1e-5
            hair_mask_sum = hair_mask.sum() + 1e-5

            render_out = self.renderer.render(rays_o, rays_d, radii, near, far,
                                              background_rgb=background_rgb,
                                              cos_anneal_ratio=self.get_cos_anneal_ratio())

            color_fine = render_out['color_fine']
            s_val_head = render_out['s_val_head']
            cdf_head = render_out['cdf_head']
            gradient_error_head = render_out['gradient_error_head']
            weight_max = render_out['weight_max']
            weight_sum = render_out['weight_sum']
        
            s_val_hair = render_out['s_val_hair']
            sampled_orient_fine = render_out['sampled_orient_fine']
            pts_fine = render_out['pts_fine']

            weights_hair = render_out['weights_hair']
            weight_hair_max = render_out['weight_hair_max']
            weight_hair_sum = render_out['weight_hair_sum']

            cdf_hair = gradient_error_hair = None

            cdf_hair = render_out['cdf_hair']
            gradient_error_hair = render_out['gradient_error_hair']

            # Loss
            color_error = (color_fine - true_rgb) * mask
            color_fine_loss = F.l1_loss(color_error, torch.zeros_like(color_error), reduction='sum') / mask_sum
            
            color_error_hair = (color_fine - true_rgb) * hair_mask
            color_hair_fine_loss = F.l1_loss(color_error, torch.zeros_like(color_error), reduction='sum') / hair_mask_sum
            
            psnr = 20.0 * torch.log10(1.0 / (((color_fine - true_rgb)**2 * mask).sum() / (mask_sum * 3.0)).sqrt())

            eikonal_loss_head = gradient_error_head


            mask_fine_loss = F.binary_cross_entropy(weight_sum.clip(1e-3, 1.0 - 1e-3), mask)

            loss = color_fine_loss +\
                   eikonal_loss_head * self.igr_weight +\
                   mask_fine_loss * self.mask_weight

            
            orient_fine_loss = orient_coarse_loss = 0.0

            if self.orient_weight:

                orient_angle_fine = project_orient_to_camera(orient_3d=sampled_orient_fine, org_3d=pts_fine, cam_intr=cam_intr, cam_extr=cam_extr)

                # Get masks for orientation loss
                orient_mask = hair_mask
                if self.conf['train']['orient_use_conf']:
                    orient_mask = orient_mask * orient_conf                    

                orient_mask_sum = orient_mask.sum() + 1e-5
                orient_fine_error = torch.min((orient_angle_fine - orient_angle).abs(),
                                              torch.min((orient_angle_fine - orient_angle - np.pi).abs(), 
                                                        (orient_angle_fine - orient_angle + np.pi).abs()))
                orient_fine_error = orient_fine_error * orient_mask 
                orient_fine_loss = orient_fine_error.mean() * (orient_fine_error.shape[0] / orient_mask_sum)

            hair_mask_fine_loss = F.binary_cross_entropy(weight_hair_sum.clip(1e-3, 1.0 - 1e-3), hair_mask)


            loss += hair_mask_fine_loss * self.mask_weight

            if self.orient_weight:
                loss += orient_fine_loss * self.orient_weight

            if self.head_prior_reg_weight:
                loss += self.head_prior_reg_weight * render_out['head_prior_inside_points'].mean()
            if self.head_prior_off_sdf_weight:
                loss += self.head_prior_off_sdf_weight * render_out['head_prior_off_sdf']
            if self.head_prior_normal_weight:
                loss += self.head_prior_normal_weight * render_out['head_prior_normal']
            if self.head_prior_sdf_weight:
                loss += self.head_prior_sdf_weight * render_out['head_prior_sdf']

            eikonal_loss_hair = gradient_error_hair
            loss += eikonal_loss_hair * self.igr_weight

            
            self.optimizer.zero_grad()
            if self.camera_model is not None and self.iter_step > self.tune_cameras_start and self.iter_step < self.tune_cameras_end:
                self.optimizer_camera.zero_grad()
            loss.backward()
            self.optimizer.step()
            if self.camera_model is not None and self.iter_step > self.tune_cameras_start and self.iter_step < self.tune_cameras_end:
                self.optimizer_camera.step()
            
            losses = {}

            self.iter_step += 1
            
            self.writer.add_scalar('Loss/hair_mask_fine_loss', hair_mask_fine_loss, self.iter_step)
            if self.orient_weight:
                self.writer.add_scalar('Loss/orient_loss', orient_fine_loss, self.iter_step)
            self.writer.add_scalar('Statistics/hair_weight_max', (weight_hair_max * hair_mask).sum() / hair_mask_sum, self.iter_step)

            if self.head_prior_reg_weight:
                self.writer.add_scalar('Loss/head_prior_reg_loss', render_out['head_prior_inside_points'].mean(), self.iter_step)
                self.writer.add_scalar('Loss/num_inside_points', render_out['num_inside_points'], self.iter_step)
            if self.head_prior_off_sdf_weight:
                self.writer.add_scalar('Loss/head_prior_off_sdf', render_out['head_prior_off_sdf'], self.iter_step)
            if self.head_prior_normal_weight:
                self.writer.add_scalar('Loss/head_prior_normal', render_out['head_prior_normal'], self.iter_step)
            if self.head_prior_sdf_weight:
                self.writer.add_scalar('Loss/head_prior_sdf', render_out['head_prior_sdf'], self.iter_step)

            self.writer.add_scalar('Statistics/s_val_hair', s_val_hair.mean(), self.iter_step)

            self.writer.add_scalar('Loss/eikonal_loss_hair', eikonal_loss_hair, self.iter_step)
            self.writer.add_scalar('Statistics/cdf_hair', (cdf_hair[:, :1] * hair_mask).sum() / hair_mask_sum, self.iter_step)

            self.writer.add_scalar('Loss/loss', loss, self.iter_step)
            self.writer.add_scalar('Loss/color_fine_loss', color_fine_loss, self.iter_step)
            self.writer.add_scalar('Loss/color_fine_hair_loss', color_hair_fine_loss, self.iter_step)
            self.writer.add_scalar('Loss/mask_fine_loss', mask_fine_loss, self.iter_step)
            self.writer.add_scalar('Loss/eikonal_loss_head', eikonal_loss_head, self.iter_step)
            self.writer.add_scalar('Statistics/s_val_head', s_val_head.mean(), self.iter_step)
            self.writer.add_scalar('Statistics/cdf_head', (cdf_head[:, :1] * mask).sum() / mask_sum, self.iter_step)
            self.writer.add_scalar('Statistics/weight_max', (weight_max * mask).sum() / mask_sum, self.iter_step)
            self.writer.add_scalar('Statistics/psnr', psnr, self.iter_step)
            if losses is not None:
                for k, v in losses.items():
                    self.writer.add_scalar(f'Loss/{k}', v, self.iter_step)

            if self.iter_step % self.report_freq == 0:
                print(self.base_exp_dir)
                print('iter:{:8>d} loss = {} lr={}'.format(self.iter_step, loss, self.optimizer.param_groups[0]['lr']))

            if self.iter_step % self.save_freq == 0:
                self.save_checkpoint()

            if self.iter_step % self.val_freq == 0:
                self.validate_image()

            if self.iter_step % self.val_mesh_freq == 0:
                self.validate_mesh(resolution=self.val_mesh_resolution, filter=False)
                
            if self.orient_weight:
                if self.iter_step % self.val_orients_freq == 0:   
                    self.validate_orientation(resolution=self.val_orients_resolution)

            self.update_learning_rate()