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.
jasonkyuyim
commented
1 year ago
Hello again. To improve runtime efficiency (slightly), in this function, you should be able to instead set
KviaK = 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