Skip to content

Commit 427e3b8

Browse files
committed
Formatted the file
1 parent 271525b commit 427e3b8

File tree

1 file changed

+7
-3
lines changed

1 file changed

+7
-3
lines changed

py/torch_tensorrt/dynamo/runtime/_ResourceAllocator.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ class ResourceAllocationStrategy(torch.nn.Module): # type: ignore[misc]
1414
def __init__(
1515
self,
1616
compiled_module: torch.nn.Module,
17-
dynamically_allocate_resources: bool = True
17+
dynamically_allocate_resources: bool = True,
1818
) -> None:
1919
super(ResourceAllocationStrategy, self).__init__()
2020
self.compiled_module = compiled_module
@@ -24,9 +24,13 @@ def __enter__(self) -> None:
2424
print("Entering resource allocator context")
2525
for name, submodule in self.compiled_module.named_modules():
2626
if "_run_on_acc" in name:
27-
submodule.use_dynamically_allocated_resources(dynamically_allocate_resources=self.dynamically_allocate_resources)
27+
submodule.use_dynamically_allocated_resources(
28+
dynamically_allocate_resources=self.dynamically_allocate_resources
29+
)
2830

2931
def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
3032
for name, submodule in self.compiled_module.named_modules():
3133
if "_run_on_acc" in name:
32-
submodule.use_dynamically_allocated_resources(dynamically_allocate_resources=self.dynamically_allocate_resources)
34+
submodule.use_dynamically_allocated_resources(
35+
dynamically_allocate_resources=self.dynamically_allocate_resources
36+
)

0 commit comments

Comments
 (0)