Getting device errors at line all_rays_cos += cos(input_rays, target_ray) in the below function.
I checked both tensors input_rays and target_ray, both are on cuda:0. I do not know which tensor is using CPU. Any help is very much useful? Thanks!
def compute_feat_weights(src_rays, all_rays, device=None):
assert src_rays.shape[0] == all_rays.shape[0]
assert src_rays.shape[-1] == all_rays.shape[-1]
SB = src_rays.shape[0]
NS = src_rays.shape[1]
H, W = src_rays.shape[2:4]
target_ray_batch_size = all_rays.shape[1]
input_rays = src_rays.view(SB, NS, -1, 8)
cos = nn.CosineSimilarity(dim=3, eps=1e-6).to(device) # similarity func
all_rays_cos = torch.zeros(SB, NS, H*W)
for idx in range(0, target_ray_batch_size):
target_ray = all_rays[:,idx].unsqueeze(0).unsqueeze(1)
print(input_rays.device, target_ray.device)
**all_rays_cos += cos(input_rays, target_ray)**
feat_weights = torch.sum(all_rays_cos, dim=2)
norm_feat_werigts = (feat_weights - torch.mean(feat_weights)) / torch.std(feat_weights)
return norm_feat_werigts.view(SB, NS, 1)`Preformatted text`