Training-RuntimeError: number of dims don't match in permute

[lyk412]

Hi! Thanks for your awesome work! I met a RuntimeError occurs in broyden.py when training. Do you have any idea about this problem? Thanks for your any help.

Traceback (most recent call last):
  File "train.py", line 72, in main
    trainer.fit(model, datamodule=datamodule)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 608, in fit
    call._call_and_handle_interrupt(
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/call.py", line 36, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 88, in launch
    return function(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 650, in _fit_impl
    self._run(model, ckpt_path=self.ckpt_path)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1112, in _run
    results = self._run_stage()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1191, in _run_stage
    self._run_train()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1214, in _run_train
    self.fit_loop.run()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py", line 267, in advance
    self._outputs = self.epoch_loop.run(self._data_fetcher)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py", line 213, in advance
    batch_output = self.batch_loop.run(kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py", line 88, in advance
    outputs = self.optimizer_loop.run(optimizers, kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/loop.py", line 199, in run
    self.advance(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 202, in advance
    result = self._run_optimization(kwargs, self._optimizers[self.optim_progress.optimizer_position])
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 249, in _run_optimization
    self._optimizer_step(optimizer, opt_idx, kwargs.get("batch_idx", 0), closure)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 370, in _optimizer_step
    self.trainer._call_lightning_module_hook(
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1356, in _call_lightning_module_hook
    output = fn(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/core/module.py", line 1742, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 169, in step
    step_output = self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 280, in optimizer_step
    optimizer_output = super().optimizer_step(optimizer, opt_idx, closure, model, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/strategies/strategy.py", line 234, in optimizer_step
    return self.precision_plugin.optimizer_step(
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 119, in optimizer_step
    return optimizer.step(closure=closure, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/optim/optimizer.py", line 88, in wrapper
    return func(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/optim/adam.py", line 92, in step
    loss = closure()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py", line 105, in _wrap_closure
    closure_result = closure()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 149, in __call__
    self._result = self.closure(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 135, in closure
    step_output = self._step_fn()
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py", line 419, in _training_step
    training_step_output = self.trainer._call_strategy_hook("training_step", *kwargs.values())
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1494, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 351, in training_step
    return self.model(*args, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 886, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/lyk/.conda/envs/gdna/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 98, in forward
    output = self._forward_module.training_step(*inputs, **kwargs)
  File "/home/lyk/gdna/lib/gdna_model.py", line 336, in training_step
    loss = self.training_step_coarse(batch, batch_idx)
  File "/home/lyk/gdna/lib/gdna_model.py", line 247, in training_step_coarse
    outputs = self.forward(batch['pts_d'], batch['smpl_tfs'],  batch['smpl_verts'], cond, eval_mode=False, only_near_smpl=False)
  File "/home/lyk/gdna/lib/gdna_model.py", line 131, in forward
    pts_c, others = self.deformer.forward(pts_d,
  File "/home/lyk/gdna/lib/model/deformer.py", line 69, in forward
    xc_opt, others = self.__search(xd, xc_init, cond, tfs, eval_mode=eval_mode)
  File "/home/lyk/gdna/lib/model/deformer.py", line 253, in __search
    result = broyden(_func, xc_init, J_inv_init, max_steps=self.opt.max_steps)
  File "/home/lyk/gdna/lib/model/broyden.py", line 69, in broyden
    J_inv[ids_val] += u.bmm(vT)
RuntimeError: number of dims don't match in permute