[Bugfix] Fix RIFE device selection for CPU-transported videos (#2876)
Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
1 parent
c0ccbb872f
commit
b7f239810d
2 files changed
+9
-5
No files matched your search
@@ -72,7 +72,7 @@ def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch):
|
||||
assert torch.isfinite(output_video).all()
|
||||
|
||||
|
||||
def test_frame_interpolator_prefers_input_tensor_device(monkeypatch):
|
||||
def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch):
|
||||
chosen_devices = []
|
||||
model = rife_interpolator.Model().eval()
|
||||
|
||||
@@ -83,10 +83,11 @@ def test_frame_interpolator_prefers_input_tensor_device(monkeypatch):
|
||||
interpolator = rife_interpolator.FrameInterpolator()
|
||||
monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded)
|
||||
monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet)
|
||||
monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda"))
|
||||
|
||||
video = torch.zeros(1, 3, 2, 32, 32)
|
||||
output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0)
|
||||
|
||||
assert chosen_devices == [video.device]
|
||||
assert chosen_devices == [torch.device("cuda")]
|
||||
assert multiplier == 2
|
||||
assert output_video.shape == (1, 3, 3, 32, 32)
|
||||
Reference in new issue
Block a user