malloc -> mallocCheck in dataloader.h

This commit is contained in:
Andrej Karpathy 2024-06-17 18:56:24 +00:00
commit 93d1450bdd

View file

@ -109,7 +109,7 @@ void prepare_intra_shard_indices_(DataLoader *loader) {
// in case shards have different number of samples / sizes
free(loader->intra_shard_indices);
}
loader->intra_shard_indices = (int*)malloc(loader->shard_num_samples * sizeof(int));
loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int));
random_permutation_with_init(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng, 1);
}
@ -177,7 +177,7 @@ void dataloader_init(DataLoader *loader,
mt19937_state shuffle_rng;
manual_seed(&shuffle_rng, 42 + process_rank);
loader->shuffle_rng = shuffle_rng;
loader->shard_indices = (int*)malloc(loader->glob_result.gl_pathc * sizeof(int));
loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int));
for (int i = 0; i < loader->glob_result.gl_pathc; i++) {
loader->shard_indices[i] = i; // start with identity permutation
}
@ -199,9 +199,9 @@ void dataloader_init(DataLoader *loader,
// printf("DataLoader: Found %ld tokens across %zu shards\n", ntok_total, loader->glob_result.gl_pathc);
// allocate all the space we'll need
loader->buffer = (uint16_t*)malloc((B * T + 1) * sizeof(uint16_t));
loader->inputs = (int*)malloc(B * T * sizeof(int));
loader->targets = (int*)malloc(B * T * sizeof(int));
loader->buffer = (uint16_t*)mallocCheck((B * T + 1) * sizeof(uint16_t));
loader->inputs = (int*)mallocCheck(B * T * sizeof(int));
loader->targets = (int*)mallocCheck(B * T * sizeof(int));
loader->num_tokens = ntok_total;
// reset the loader, to initialize it
@ -359,11 +359,11 @@ void evalloader_init(EvalLoader *loader,
// allocate all the space we'll need
int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS;
loader->buffer = (uint16_t*)malloc(longest_example_bytes);
loader->buffer = (uint16_t*)mallocCheck(longest_example_bytes);
loader->inputs = (int*)calloc(B * T, sizeof(int));
loader->targets = (int*)calloc(B * T, sizeof(int));
loader->mask = (char*)malloc(B * T * sizeof(char));
loader->label = (int*)malloc(can_fit_examples * sizeof(int));
loader->mask = (char*)mallocCheck(B * T * sizeof(char));
loader->label = (int*)mallocCheck(can_fit_examples * sizeof(int));
// reset the loader, to initialize it
evalloader_reset(loader);