Hello again. To improve runtime efficiency (slightly), in this function, you should be able to instead set K via K = torch.arange(embed_size//2, device=indices.device). Then, you should be able to remove the two .to() calls in the remainder of this function since K (and its Tensor derivatives) will already be loaded onto the correct device.
Hello again. To improve runtime efficiency (slightly), in this function, you should be able to instead set
K
viaK = torch.arange(embed_size//2, device=indices.device)
. Then, you should be able to remove the two.to()
calls in the remainder of this function sinceK
(and its Tensor derivatives) will already be loaded onto the correct device.https://github.com/jasonkyuyim/se3_diffusion/blob/420b1fe6d3da53ff17490e5b599df2acca5039e5/model/score_network.py#L25