test_srt.py

import torch

from srt.encoder import SRTEncoder
from srt.decoder import SRTDecoder, NerfDecoder

import yaml

batch_size = 1
num_views = 2
device = torch.device("cuda:0")


encoder = SRTEncoder().to(device)
decoder = NerfDecoder().to(device)

images = torch.rand(batch_size, num_views, 3, 224, 224).to(device)
i_rays = torch.rand(batch_size, num_views, 224, 224, 3).to(device)
i_positions = torch.rand(batch_size, num_views, 3).to(device)
t_rays = torch.rand(batch_size, 224*224, 3).to(device)
t_positions = torch.rand(batch_size, 224*224, 3).to(device)

z = encoder(images, i_positions, i_rays)
recons = decoder(z, t_positions, t_rays)

loss = torch.nn.functional.mse_loss(recons, images)