Fix BertScore on GPU

[SkafteNicki]

What does this PR do?

Fixes #909

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #912 (55770a1) into master (f055686) will increase coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff          @@
##           master   #912   +/-   ##
=====================================
  Coverage      95%    95%           
=====================================
  Files         173    173           
  Lines        7344   7344           
=====================================
+ Hits         6946   6963   +17     
+ Misses        398    381   -17     

[stancld]

@SkafteNicki @Borda Won't be there any problem with GPU memory when the embeddings are going to be stored in its RAM?

[SkafteNicki]

@stancld true, i guess the we need to move to GPU at some later point in the computation?

[stancld]

Hi @SkafteNicki, I'd rather suggest the following :]. I can see that input_ids_idf wasn't enforced to move to cpu.

out = torch.einsum("blsd, bs -> blsd", out, processed_attention_mask)
        embeddings_list.append(out.cpu())

        # Calculate weighted (w.r.t. sentence length) input_ids IDF matrix
        input_ids_idf = (
            batch["input_ids_idf"] * processed_attention_mask if idf else processed_attention_mask.type(out.dtype)
        )
        input_ids_idf /= input_ids_idf.sum(-1, keepdim=True)
-         idf_scale_list.append(input_ids_idf)
+         idf_scale_list.append(input_ids_idf.cpu())

[stancld]

@SkafteNicki @Borda Please, have a look at minor change I made. I believe it can be marked as ready for review :)

[Borda]

can we add a test to cover this case? 🐰

[stancld]

can we add a test to cover this case? 🐰

Should be tested thanks to a test parameter added by @SkafteNicki:

@pytest.mark.parametrize("device", ["cpu", "cuda"])

Tests fail without our code change and with device="cuda".

[justusschock]

LGTM! Haven't looked at the implementation, but I assume it relies on .to(device=device), where we should also test for no explicit device