Commit 174907b6 authored by Peter Eastman's avatar Peter Eastman
Browse files

Fixed error loading dataset without weights

parent ec10283f
Loading
Loading
Loading
Loading
+10 −4
Original line number Diff line number Diff line
@@ -507,8 +507,11 @@ class DiskDataset(Dataset):
            os.path.join(dataset.data_dir, row['X'])))
        y = np.array(load_from_disk(
            os.path.join(dataset.data_dir, row['y'])))
        w = np.array(load_from_disk(
            os.path.join(dataset.data_dir, row['w'])))
        w_filename = os.path.join(dataset.data_dir, row['w'])
        if os.path.exists(w_filename):
            w = np.array(load_from_disk(w_filename))
        else:
            w = np.ones(y.shape)
        ids = np.array(load_from_disk(
            os.path.join(dataset.data_dir, row['ids'])), dtype=object)
        yield (X, y, w, ids)
@@ -725,8 +728,11 @@ class DiskDataset(Dataset):
        os.path.join(self.data_dir, row['X'])))
    y = np.array(load_from_disk(
        os.path.join(self.data_dir, row['y'])))
    w = np.array(load_from_disk(
        os.path.join(self.data_dir, row['w'])))
    w_filename = os.path.join(self.data_dir, row['w'])
    if os.path.exists(w_filename):
        w = np.array(load_from_disk(w_filename))
    else:
        w = np.ones(y.shape)
    ids = np.array(load_from_disk(
        os.path.join(self.data_dir, row['ids'])), dtype=object)
    return (X, y, w, ids)