Skip to content

Commit 613042e

Browse files
authored
Handle optional to_numpy
1 parent 9471f17 commit 613042e

File tree

1 file changed

+8
-2
lines changed

1 file changed

+8
-2
lines changed

tests/test_implicit.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,11 @@ def test_model(self):
2828
model.fit(counts, show_progress=False)
2929
rows, cols = model.item_factors, model.user_factors
3030

31-
assert not np.isnan(np.sum(cols.to_numpy()))
32-
assert not np.isnan(np.sum(rows.to_numpy()))
31+
assert not np.isnan(np.sum(tonumpy(cols))
32+
assert not np.isnan(np.sum(tonumpy(rows))
33+
34+
35+
def tonumpy(x):
36+
if hasattr(x, 'to_numpy'):
37+
return x.to_numpy()
38+
return x

0 commit comments

Comments
 (0)