Replaceinference_mode() with no_grad() to resolve training (#559)

Author: Borda

rfdetr/detr.py

@@ -310,7 +310,7 @@ def predict(
                                      "Alternatively, you can recompile the optimized model for a different batch size "
                                      "by calling model.optimize_for_inference(batch_size=<new_batch_size>).")
 
-        with torch.inference_mode():
+        with torch.no_grad():
             if self._is_optimized_for_inference:
                 predictions = self.model.inference_model(batch_tensor.to(dtype=self._optimized_dtype))
             else:

rfdetr/engine.py

@@ -113,7 +113,7 @@ def train_one_epoch(
             scales = compute_multi_scale_scales(args.resolution, args.expanded_scales, args.patch_size, args.num_windows)
             random.seed(it)
             scale = random.choice(scales)
-            with torch.inference_mode():
+            with torch.no_grad():
                 samples.tensors = F.interpolate(samples.tensors, size=scale, mode='bilinear', align_corners=False)
                 samples.mask = F.interpolate(samples.mask.unsqueeze(1).float(), size=scale, mode='nearest').squeeze(1).bool()
 

rfdetr/main.py

@@ -368,7 +368,7 @@ def lr_lambda(current_step: int):
 
                         utils.save_on_master(weights, checkpoint_path)
 
-            with torch.inference_mode():
+            with torch.no_grad():
                 test_stats, coco_evaluator = evaluate(
                     model, criterion, postprocess, data_loader_val, base_ds, device, args=args
                 )