| 451 | yield name, b |
| 452 | |
| 453 | def change_tensor_cuda(): |
| 454 | sync_print('override tensor.cuda() to preserve task_specific flag') |
| 455 | # change .cuda of Tensor |
| 456 | ori_tensor_cuda = torch.Tensor.cuda |
| 457 | torch.Tensor.ori_cuda = ori_tensor_cuda |
| 458 | def new_cuda(self, *args, **kwargs): # could be written as decorator I guess... |
| 459 | cuda_t = self.ori_cuda(*args, **kwargs) |
| 460 | if hasattr(self, 'task_specific'): |
| 461 | cuda_t.task_specific = self.task_specific |
| 462 | if hasattr(self, 'backbone_specific'): |
| 463 | cuda_t.backbone_specific = self.backbone_specific |
| 464 | if hasattr(self, 'neck_specific'): |
| 465 | cuda_t.neck_specific = self.neck_specific |
| 466 | if hasattr(self, 'decoder_specific'): |
| 467 | cuda_t.decoder_specific = self.decoder_specific |
| 468 | return cuda_t |
| 469 | torch.Tensor.cuda = new_cuda |
| 470 | |
| 471 | def add_task_specific(m, task_specific): |
| 472 | for name, param in m.named_parameters(): |