henry000 commited on
Commit
70a7f92
·
1 Parent(s): 44534b9

:white_check_mark: [Pass] the test with fixing model location

Browse files
tests/test_model/test_yolo.py CHANGED
@@ -37,8 +37,8 @@ def cfg() -> Config:
37
  @pytest.fixture
38
  def model(cfg: Config):
39
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
40
- model = create_model(cfg.model, weight_path=None)
41
- return model.to(device)
42
 
43
 
44
  def test_model_basic_status(model):
 
37
  @pytest.fixture
38
  def model(cfg: Config):
39
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
40
+ model = create_model(cfg.model, weight_path=None, device=device)
41
+ return model
42
 
43
 
44
  def test_model_basic_status(model):
tests/test_utils/test_loss.py CHANGED
@@ -24,8 +24,8 @@ def cfg() -> Config:
24
  @pytest.fixture
25
  def model(cfg: Config):
26
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27
- model = create_model(cfg.model, weight_path=None)
28
- return model.to(device)
29
 
30
 
31
  @pytest.fixture
 
24
  @pytest.fixture
25
  def model(cfg: Config):
26
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27
+ model = create_model(cfg.model, weight_path=None, device=device)
28
+ return model
29
 
30
 
31
  @pytest.fixture