add view to input

2023-01-10 14:02:11 +00:00
parent 337ff40c89
commit c7ba2f1b54
3 changed files with 31 additions and 7 deletions
@@ -23,7 +23,7 @@ def train_step(batch):
    model.train() # train mode is on by default after construction

    # 1. Forward pass
-    y_pred = model(batch['rgb_image'])
+    y_pred = model([batch['rgb_image'], batch['view']])

    # 2. Calculate loss
    loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
@@ -45,7 +45,7 @@ def valid_step(batch):
    model.eval() # put the model in evaluation mode for testing (inference)
    # 1. Forward pass
    with torch.inference_mode():
-        pred = model(batch['rgb_image'])
+        pred = model([batch['rgb_image'], batch['view']])
    
        # 2. Calculate the loss
        loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])