Merge pull request #43 from leobavila/fix-amp-get-pytorch-version

Update amp.py

Author: cadurosar

splade/tasks/amp.py

@@ -2,7 +2,7 @@
 
 # inspired from Colbert repo: https://github.com/stanford-futuredata/ColBERT
 
-PyTorch_over_1_6 = float((torch.__version__.split('.')[1])) >= 6 and float((torch.__version__.split('.')[0])) >= 1
+PyTorch_over_1_6 = float(".".join([torch.__version__.split('.')[0], torch.__version__.split('.')[1]])) >= 1.6
 
 
 # replace this with  contextlib.nullcontext if python >3.7