def test_get_item(): diag_lv = DiagLazyVariable(Variable(diag)) diag_ev = diag_lv.evaluate() assert torch.equal(diag_lv[0:2].evaluate().data, diag_ev[0:2].data)