allocate memory on batch generation (#126)

mhubii · mhubii · commit 5cd9283dd028 · 2026-02-04T15:28:12.000Z
diff --git a/roboreg/util/transform.py b/roboreg/util/transform.py
@@ -64,7 +64,7 @@ def generate_ht_optical(
 ) -> torch.Tensor:
     ht_optical = torch.zeros(4, 4, dtype=dtype, device=device)
     if batch_size is not None:
-        ht_optical = ht_optical.unsqueeze(0).expand(batch_size, -1, -1)
+        ht_optical = ht_optical.unsqueeze(0).repeat(batch_size, 1, 1)
     ht_optical[..., 0, 2] = (
         1.0  # OpenCV-oriented optical frame, in quaternions: [0.5, -0.5, 0.5, -0.5] (w, x, y, z)
     )

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ def generate_ht_optical(`
`64`	`64`	`) -> torch.Tensor:`
`65`	`65`	`ht_optical = torch.zeros(4, 4, dtype=dtype, device=device)`
`66`	`66`	`if batch_size is not None:`
`67`		`- ht_optical = ht_optical.unsqueeze(0).expand(batch_size, -1, -1)`
	`67`	`+ ht_optical = ht_optical.unsqueeze(0).repeat(batch_size, 1, 1)`
`68`	`68`	`ht_optical[..., 0, 2] = (`
`69`	`69`	`1.0 # OpenCV-oriented optical frame, in quaternions: [0.5, -0.5, 0.5, -0.5] (w, x, y, z)`
`70`	`70`	`)`