[Pytorch] RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling

CUBLAS_STATUS_NOT_INITIALIZED 에러는 nn.Embedding에 원래 값보다 큰 값이 들어갔을 때 발생하는 에러다. 예를 들어, 사용자의 수는 100명이고, 제품의 수는 50개라고 하자. 그렇다면 아래와 같은 코드로 nn.Embedding을 생성할 수 있을 것이다.

class Model(nn.Mudule):
	def __init__(self, num_users, num_items, latent_dim):
    	self.num_users = num_users
        self.num_items = num_items
        self.latent_dim = latent_dim
        
        self.user_embedding = nn.Embedding(self.num_users, self.latent_dim)
        self.item_embedding = nn.Embedding(self.num_items, self.latent_dim)
        
        self.fc_layer = nn.Linear(self.latent_dim, 1)
        self.sigmoid = nn.Sigmoid()
    
    def forward(self, uid, iid):
    	self.user_emb = self.user_embedding(uid)
        self.item_emb = self.item_embedding(iid)
        
        multiply = torch.mul(self.user_emb, self.item_emb)
        output = self.fc_layer(multiply)
        output = self.sigmoid(output)
        return output

 

사용자의 ID는 0부터 99까지 총 100개의 ID 값을 가지게 될 것이고, 제품의 ID는 0부터 49까지 총 50개의 ID 값을 가지게 될 것이다. 이때 forward 부분에 입력으로 사용자의 ID와 제품의 ID가 들어오는데, 100 이상의 ID 값이 들어올 경우 위와 같은 에러를 출력한다. 이 문제를 해결하기 위해서는 사용자 ID와 제품 ID의 max값을 설정하는 등의 처리를 하면 된다.