added activation for features

2023-01-10 13:37:31 +00:00
parent bca2e0bd2c
commit 337ff40c89
3 changed files with 11 additions and 6 deletions
@@ -70,7 +70,7 @@ if __name__ == '__main__':
    train_accumulator = Accumulator()
    valid_accumulator = Accumulator()

-    model = DepthEstimation().to('cuda')
+    model = DepthEstimation(device=device)

    loss_fn = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)