R-GCN run time error during evaluation

Hi,

I get the following error when running R-GCN with FB15k-237. It happens during evaluation. I reduced the eval batch size but still after 2000 epochs it fails with this run time error.

File “/home/saatviga_sudhahar/workspace/dgl/examples/pytorch/rgcn/utils.py”, line 210, in calc_mrr
ranks_s = perturb_and_get_rank(embedding, w, o, r, s, test_size, eval_bz)
File “/home/saatviga_sudhahar/workspace/dgl/examples/pytorch/rgcn/utils.py”, line 193, in perturb_and_get_rank
out_prod = torch.bmm(emb_ar, emb_c) # size D x E x V
RuntimeError: [enforce fail at CPUAllocator.cpp:64] . DefaultCPUAllocator: can’t allocate memory: you tried to allocate 3722496000 bytes. Error code 12 (Cannot allocate memory)

Has anyone faced the same issue?

Did you train your model on CPU or GPU?