(coords, colors, model)
| 95 | |
| 96 | |
| 97 | def train(coords, colors, model): |
| 98 | # Measure time |
| 99 | timer = Timer() |
| 100 | for i in range(5): |
| 101 | # Feed-forward pass and get the prediction |
| 102 | timer.tic() |
| 103 | sinput = ME.SparseTensor( |
| 104 | colors, |
| 105 | coords, |
| 106 | device=device, |
| 107 | allocator_type=ME.GPUMemoryAllocatorType.PYTORCH, |
| 108 | ) |
| 109 | logits = model(sinput) |
| 110 | logits.F.sum().backward() |
| 111 | timer.toc() |
| 112 | return timer.min_time, len(logits) |
| 113 | |
| 114 | |
| 115 | def test_network(coords, feats, model, batch_sizes, forward_only=True): |