Skip to content

Commit

Permalink
Update backbone example for v2.1.0 (#224)
Browse files Browse the repository at this point in the history
Co-authored-by: ys-2020 <shangy@mit.edu>
  • Loading branch information
ys-2020 and ys-2020 committed Jul 28, 2023
1 parent 897096a commit b55506a
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions examples/backbones.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
@torch.no_grad()
def main() -> None:
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
from torchsparse.nn import functional as F
F.set_kmap_mode('hashmap')

for backbone in [SparseResNet21D, SparseResUNet42]:
print(f'{backbone.__name__}:')
Expand All @@ -23,8 +25,8 @@ def main() -> None:
pcs -= np.min(pcs, axis=0, keepdims=True)
pcs, indices = sparse_quantize(pcs, voxel_size, return_index=True)
coords = np.zeros((pcs.shape[0], 4))
coords[:, :3] = pcs[:, :3]
coords[:, -1] = 0
coords[:, 1:4] = pcs[:, :3]
coords[:, 0] = 0
coords = torch.as_tensor(coords, dtype=torch.int)
feats = torch.as_tensor(feats[indices], dtype=torch.float)
input = SparseTensor(coords=coords, feats=feats).to(device)
Expand Down

0 comments on commit b55506a

Please sign in to comment.