style

Author: ssheorey

ml3d/torch/models/point_transformer.py

@@ -730,8 +730,9 @@ def knn_batch(points,
                      return_distances=True)
     if return_distances:
         return ans.neighbors_index.reshape(
-            -1, k).long().to(device), ans.neighbors_distance.reshape(-1, k).to(
-                device)
+            -1,
+            k).long().to(device), ans.neighbors_distance.reshape(-1,
+                                                                 k).to(device)
     else:
         return ans.neighbors_index.reshape(-1, k).long().to(device)
 

ml3d/torch/pipelines/base_pipeline.py

@@ -72,8 +72,7 @@ def __init__(self,
             if distributed:
                 raise NotImplementedError(
                     "Distributed training for XPU is not supported yet.")
-            self.device = torch.device(
-                device if ':' in str(device) else 'xpu')
+            self.device = torch.device(device if ':' in str(device) else 'xpu')
         elif not torch.cuda.is_available():
             if distributed:
                 raise NotImplementedError(
@@ -83,8 +82,8 @@ def __init__(self,
             if distributed:
                 self.device = torch.device(device)
                 print(f"Rank : {self.rank} using device : {self.device}")
-                device_index = (self.device.index if self.device.index is not None
-                                else 0)
+                device_index = (self.device.index
+                                if self.device.index is not None else 0)
                 if hasattr(torch, 'accelerator'):
                     torch.accelerator.set_device_index(device_index)
                 elif self.device.type == 'cuda':

ml3d/torch/pipelines/object_detection.py

@@ -336,8 +336,8 @@ def run_train(self):
         # wrap model for multiple GPU
         if self.distributed:
             model.to(self.device)
-            device_index = (self.device.index if self.device.index is not None
-                            else 0)
+            device_index = (self.device.index
+                            if self.device.index is not None else 0)
             model = torch.nn.parallel.DistributedDataParallel(
                 model, device_ids=[device_index])
             model.get_loss = model.module.get_loss

tests/test_models_torch.py

@@ -47,8 +47,7 @@
 import pytest
 import torch
 
-from torch_backend_parity import (assert_cpu_accelerator_parity,
-                                  move_to_device)
+from torch_backend_parity import (assert_cpu_accelerator_parity, move_to_device)
 
 if 'PATH_TO_OPEN3D_ML' in os.environ.keys():
     base = os.environ['PATH_TO_OPEN3D_ML']
@@ -92,8 +91,8 @@
 
 def _bev_box3d():
     """BEVBox3D lives in Open3D-ML (ml3d), not in open3d.ml.datasets."""
-    ml_root = (os.environ.get('OPEN3D_ML_ROOT')
-               or os.environ.get('PATH_TO_OPEN3D_ML') or base)
+    ml_root = (os.environ.get('OPEN3D_ML_ROOT') or
+               os.environ.get('PATH_TO_OPEN3D_ML') or base)
     if ml_root not in sys.path:
         sys.path.insert(0, ml_root)
     from ml3d.datasets.utils import BEVBox3D
@@ -115,8 +114,7 @@ def _pointrcnn_ctor_kwargs(mode):
             npoints=[256, 64, 16, 4],
             radius=[[0.1, 0.5], [0.5, 1.0], [1.0, 2.0], [2.0, 4.0]],
             nsample=[[4, 8], [4, 8], [4, 8], [4, 8]],
-            mlps=[[[8, 8, 16], [16, 16, 32]],
-                  [[32, 32, 64], [32, 48, 64]],
+            mlps=[[[8, 8, 16], [16, 16, 32]], [[32, 32, 64], [32, 48, 64]],
                   [[64, 98, 128], [64, 98, 128]],
                   [[128, 128, 256], [128, 192, 256]]],
         ),
@@ -500,8 +498,8 @@ def make_net(device=torch.device('cpu')):
     def run(device):
         model = make_net(device)
         model.load_state_dict(state)
-        batcher = ml3d.dataloaders.ConcatBatcher(
-            str(device), model='PointTransformer')
+        batcher = ml3d.dataloaders.ConcatBatcher(str(device),
+                                                 model='PointTransformer')
         batch = batcher.collate_fn([sample])
         batch['data'].to(device)
         out = model(batch['data'])

tests/torch_backend_parity.py

@@ -7,6 +7,7 @@
 RTOL = {"cuda": 1e-4, "xpu": 1e-3}
 ATOL = {"cuda": 1e-5, "xpu": 1e-4}
 
+
 def move_to_device(obj, device):
     """Recursively move tensors in nested batch structures to ``device``."""
     if torch.is_tensor(obj):