Make sure test tensors are in inference mode

Author: guill

benchmarks/benchmark.py

@@ -254,7 +254,8 @@ def example_entrypoint():
                         else:
                             print(f"  Creating {name} tensor {size}...")
 
-                        tensor = torch.randn(*size)
+                        with torch.inference_mode():
+                            tensor = torch.randn(*size)
                         test_data.append((f"{name}_cpu", tensor))
 
                         size_gb = (tensor.numel() * 4) / (1024**3)
@@ -266,11 +267,13 @@ def example_entrypoint():
                                 # Skip GPU for very large tensors to avoid OOM
                                 if name == "image_8k" or name == "model_6gb":
                                     print(f"    Creating GPU version of {name} (may use significant VRAM)...")
-                                    gpu_tensor = tensor.cuda()
+                                    with torch.inference_mode():
+                                        gpu_tensor = tensor.cuda()
                                     test_data.append((f"{name}_gpu", gpu_tensor))
                                     print("    GPU tensor created successfully")
                                 else:
-                                    gpu_tensor = tensor.cuda()
+                                    with torch.inference_mode():
+                                        gpu_tensor = tensor.cuda()
                                     test_data.append((f"{name}_gpu", gpu_tensor))
                                     print("    GPU tensor created successfully")
                             except RuntimeError as gpu_e:

benchmarks/memory_benchmark.py

@@ -309,10 +309,11 @@ async def run_scaling_test(
 
             # Create test tensor
             print(f"Creating test tensor {test_tensor_size}...")
-            if use_cuda and CUDA_AVAILABLE:
-                test_tensor = torch.randn(*test_tensor_size, device="cuda")
-            else:
-                test_tensor = torch.randn(*test_tensor_size)
+            with torch.inference_mode():
+                if use_cuda and CUDA_AVAILABLE:
+                    test_tensor = torch.randn(*test_tensor_size, device="cuda")
+                else:
+                    test_tensor = torch.randn(*test_tensor_size)
 
             tensor_size_mb = test_tensor.element_size() * test_tensor.numel() / (1024 * 1024)
             print(f"Tensor size: {tensor_size_mb:.1f} MB on {test_tensor.device}")
@@ -479,7 +480,10 @@ async def run_large_tensor_sharing_test(
                 side = int(num_elements**0.5)
 
                 print(f"Creating {tensor_gb}GB tensor ({side}x{side}) on {device_name}...")
-                large_tensor = torch.randn(side, side, device="cuda") if use_cuda else torch.randn(side, side)
+                with torch.inference_mode():
+                    large_tensor = (
+                        torch.randn(side, side, device="cuda") if use_cuda else torch.randn(side, side)
+                    )
                 actual_size_mb = large_tensor.element_size() * large_tensor.numel() / (1024 * 1024)
                 print(f"Actual tensor size: {actual_size_mb:.1f} MB on {large_tensor.device}")
 

benchmarks/simple_benchmark.py

@@ -122,20 +122,22 @@ class CustomConfig(TypedDict):
             for name, size in tensor_specs:
                 try:
                     print(f"  Creating {name} tensor {size}...")
-                    tensor = torch.randn(*size)
+                    with torch.inference_mode():
+                        tensor = torch.randn(*size)
                     test_data.append((name, tensor))
                     print(f"    {name} created successfully ({tensor.numel() * 4 / (1024**3):.2f}GB)")
                 except RuntimeError as e:
                     print(f"    Skipping {name}: {e}")
 
             if include_large_tensors:
                 print("  Including very large tensors (this will use significant memory)...")
-                test_data.extend(
-                    [
-                        ("huge_tensor", torch.randn(4096, 4096)),  # ~64MB
-                        ("image_4k", torch.randn(3, 4096, 4096)),  # ~200MB (4K RGB image)
-                    ]
-                )
+                with torch.inference_mode():
+                    test_data.extend(
+                        [
+                            ("huge_tensor", torch.randn(4096, 4096)),  # ~64MB
+                            ("image_4k", torch.randn(3, 4096, 4096)),  # ~200MB (4K RGB image)
+                        ]
+                    )
                 # 8K image would be ~800MB, only add if explicitly requested
                 print("  (8K image tensor skipped - would use ~800MB)")
             else:

tests/test_benchmarks.py

@@ -394,7 +394,8 @@ async def test_torch_tensor_benchmarks(self):
         print("=" * 60)
 
         # Small tensor (CPU)
-        small_tensor_cpu = torch.randn(100, 100)  # ~40KB
+        with torch.inference_mode():
+            small_tensor_cpu = torch.randn(100, 100)  # ~40KB
 
         await self.runner.run_benchmark(
             "Small Tensor CPU - Local Baseline",
@@ -406,7 +407,8 @@ async def test_torch_tensor_benchmarks(self):
         )
 
         # Large tensor (CPU)
-        large_tensor_cpu = torch.randn(1024, 1024)  # ~4MB
+        with torch.inference_mode():
+            large_tensor_cpu = torch.randn(1024, 1024)  # ~4MB
 
         await self.runner.run_benchmark(
             "Large Tensor CPU - RPC Call",
@@ -416,8 +418,9 @@ async def test_torch_tensor_benchmarks(self):
 
         # GPU tests if available
         if CUDA_AVAILABLE:
-            small_tensor_gpu = small_tensor_cpu.cuda()
-            large_tensor_gpu = large_tensor_cpu.cuda()
+            with torch.inference_mode():
+                small_tensor_gpu = small_tensor_cpu.cuda()
+                large_tensor_gpu = large_tensor_cpu.cuda()
 
             await self.runner.run_benchmark(
                 "Small Tensor GPU - RPC Call",

tests/test_torch_tensor_integration.py

@@ -267,7 +267,8 @@ def example_entrypoint() -> ExampleExtension:
             # Test 1: Simple CPU tensor
             import torch
 
-            cpu_tensor = torch.randn(3, 4)
+            with torch.inference_mode():
+                cpu_tensor = torch.randn(3, 4)
 
             # Call extension method
             result_tensor = await extension.do_stuff({"operation": "process_tensor", "tensor": cpu_tensor})
@@ -284,7 +285,8 @@ def example_entrypoint() -> ExampleExtension:
             assert tensor_info["is_cuda"] is False
 
             # Test 2: Multiple tensors
-            tensors = [torch.ones(2, 2), torch.zeros(2, 2), torch.eye(2)]
+            with torch.inference_mode():
+                tensors = [torch.ones(2, 2), torch.zeros(2, 2), torch.eye(2)]
             stacked_result = await extension.do_stuff(
                 {"operation": "test_multiple_tensors", "tensors": tensors}
             )
@@ -415,7 +417,8 @@ def example_entrypoint() -> ExampleExtension:
             import torch
 
             # Test 1: Basic tensor processing
-            input_tensor = torch.randn(4, 5)
+            with torch.inference_mode():
+                input_tensor = torch.randn(4, 5)
             normalized = await extension.do_stuff(
                 {"operation": "process_tensor_isolated", "tensor": input_tensor}
             )
@@ -428,11 +431,12 @@ def example_entrypoint() -> ExampleExtension:
             assert abs(norm_info["output_std"] - 1.0) < 1e-6  # Should be close to 1
 
             # Test 2: Different dtypes
-            tensors_dict = {
-                "float32": torch.randn(2, 3),
-                "int64": torch.randint(0, 10, (2, 3)),
-                "bool": torch.tensor([[True, False], [False, True]]),
-            }
+            with torch.inference_mode():
+                tensors_dict = {
+                    "float32": torch.randn(2, 3),
+                    "int64": torch.randint(0, 10, (2, 3)),
+                    "bool": torch.tensor([[True, False], [False, True]]),
+                }
 
             dtype_results = await extension.do_stuff(
                 {"operation": "test_different_dtypes", "tensors_dict": tensors_dict}
@@ -540,7 +544,8 @@ def example_entrypoint() -> ExampleExtension:
             import torch
 
             # Test 1: GPU tensor operations
-            gpu_tensor = torch.randn(5, 5).cuda()
+            with torch.inference_mode():
+                gpu_tensor = torch.randn(5, 5).cuda()
             gpu_result = await extension.do_stuff({"operation": "process_gpu_tensor", "tensor": gpu_tensor})
 
             assert isinstance(gpu_result, torch.Tensor)
@@ -552,7 +557,8 @@ def example_entrypoint() -> ExampleExtension:
             assert "cuda" in gpu_info["device"]
 
             # Test 2: CPU to GPU transfer
-            cpu_tensor = torch.ones(3, 3)
+            with torch.inference_mode():
+                cpu_tensor = torch.ones(3, 3)
             transferred_result = await extension.do_stuff(
                 {"operation": "transfer_between_devices", "tensor": cpu_tensor}
             )
@@ -637,7 +643,8 @@ def example_entrypoint() -> ExampleExtension:
             import torch
 
             # Test GPU operations
-            gpu_tensor = torch.randn(4, 4).cuda()
+            with torch.inference_mode():
+                gpu_tensor = torch.randn(4, 4).cuda()
             squared_result = await extension.do_stuff(
                 {"operation": "process_gpu_operations", "tensor": gpu_tensor}
             )