Junwu_Weng:

thX = torch.squeeze(self.r_para[b_idx, t_idx, 0]) * PI


If thX is a one-dimensional tensor containing one element you can just do thX[0].
With regards to your second question, I’m not sure what you’re trying to do with your code. If you could explain a little more what you’re doing with the two for loops we might be able to find a better way