From 8e6dc2878535b37c7bbcaa086825b9da0c44d629 Mon Sep 17 00:00:00 2001 From: Shahzad Malik Muzaffar Date: Wed, 13 May 2026 08:05:47 +0200 Subject: [PATCH] Fix for testModelJit unit test --- PhysicsTools/PyTorch/interface/Model.h | 2 +- PhysicsTools/PyTorch/test/testModelJit.cc | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/PhysicsTools/PyTorch/interface/Model.h b/PhysicsTools/PyTorch/interface/Model.h index cf390eff8f47f..c0efd2df4fab0 100644 --- a/PhysicsTools/PyTorch/interface/Model.h +++ b/PhysicsTools/PyTorch/interface/Model.h @@ -28,7 +28,7 @@ namespace cms::torch { if (dev == device_) return; - assert(!is_frozen_ && "Model is frozen, cannot be moved to another device!"); + TORCH_CHECK(!is_frozen_ && "Model is frozen, cannot be moved to another device!"); model_.to(dev, non_blocking); device_ = dev; if (auto_freeze_) { diff --git a/PhysicsTools/PyTorch/test/testModelJit.cc b/PhysicsTools/PyTorch/test/testModelJit.cc index 52b296767951b..c3b459b0743f0 100644 --- a/PhysicsTools/PyTorch/test/testModelJit.cc +++ b/PhysicsTools/PyTorch/test/testModelJit.cc @@ -78,7 +78,7 @@ namespace torchtest { auto m_path = edm::FileInPath(modelPath).fullPath(); forEachCudaDevice([&](auto dev) { - auto m = cms::torch::Model(m_path); + auto m = cms::torch::Model(m_path, false); m.to(dev); CPPUNIT_ASSERT_EQUAL(dev, m.device()); @@ -86,6 +86,14 @@ namespace torchtest { m.to(::torch::kCPU); CPPUNIT_ASSERT_EQUAL(::torch::kCPU, m.device().type()); }); + + forEachCudaDevice([&](auto dev) { + auto m = cms::torch::Model(m_path); + m.to(dev); + + CPPUNIT_ASSERT_EQUAL(dev, m.device()); + CPPUNIT_ASSERT_THROW(m.to(::torch::kCPU), c10::Error); + }); } void TestModelJIT::testToDevice_NonBlocking() {